updates

agent/internal/sync/sync.go

@@ -162,9 +162,28 @@ func connectAndHandleStream(ctx context.Context, cfg *config.Config) error {
 		if cmd.GenerateKey != nil {
 			go handleGenerateKey(cfg, cmd)
 		}
+		if cmd.DeleteKey != nil {
+			go handleDeleteKey(cmd)
+		}
 	}
 }
 
+func handleDeleteKey(cmd *pb.ServerCommand) {
+	label := cmd.DeleteKey.Label
+	keyPath := fmt.Sprintf("/root/.ssh/keymanager_%s", strings.ReplaceAll(label, " ", "_"))
+
+	if err := keys.RemoveSSHIdentity(keyPath); err != nil {
+		log.Printf("remove ssh identity failed (cmd=%s): %v", cmd.CommandId, err)
+	}
+
+	for _, path := range []string{keyPath, keyPath + ".pub"} {
+		if err := os.Remove(path); err != nil && !os.IsNotExist(err) {
+			log.Printf("delete key file %s (cmd=%s): %v", path, cmd.CommandId, err)
+		}
+	}
+	log.Printf("deleted local key files for %q (cmd=%s)", label, cmd.CommandId)
+}
+
 func handleGenerateKey(cfg *config.Config, cmd *pb.ServerCommand) {
 	g := cmd.GenerateKey
 	label := g.Label
@@ -182,6 +201,12 @@ func handleGenerateKey(cfg *config.Config, cmd *pb.ServerCommand) {
 		return
 	}
 
+	privKeyData, err := os.ReadFile(keyPath)
+	if err != nil {
+		log.Printf("read private key failed (cmd=%s): %v", cmd.CommandId, err)
+		return
+	}
+
 	client, err := grpcclient.New(cfg.ServerURL, cfg.TLS)
 	if err != nil {
 		log.Printf("dial for key upload failed (cmd=%s): %v", cmd.CommandId, err)
@@ -189,11 +214,15 @@ func handleGenerateKey(cfg *config.Config, cmd *pb.ServerCommand) {
 	}
 	defer client.Close()
 
-	keyID, err := client.UploadGeneratedKey(cfg.ServerID, cfg.AgentToken, pubKey, label)
+	keyID, err := client.UploadGeneratedKey(cfg.ServerID, cfg.AgentToken, pubKey, string(privKeyData), label)
 	if err != nil {
 		log.Printf("key upload failed (cmd=%s): %v", cmd.CommandId, err)
 		return
 	}
+
+	if err := keys.AddSSHIdentity(keyPath); err != nil {
+		log.Printf("add ssh identity failed (cmd=%s): %v", cmd.CommandId, err)
+	}
 	log.Printf("generated and uploaded key %q (key_id=%s, cmd=%s)", label, keyID, cmd.CommandId)
 }
 
@@ -226,10 +255,19 @@ func GenerateAndUpload(cfg *config.Config, label string) error {
 		return err
 	}
 
-	keyID, err := client.UploadGeneratedKey(cfg.ServerID, cfg.AgentToken, pubKey, label)
+	privKeyData, err := os.ReadFile(keyPath)
+	if err != nil {
+		return fmt.Errorf("read private key: %w", err)
+	}
+
+	keyID, err := client.UploadGeneratedKey(cfg.ServerID, cfg.AgentToken, pubKey, string(privKeyData), label)
 	if err != nil {
 		return err
 	}
+
+	if err := keys.AddSSHIdentity(keyPath); err != nil {
+		log.Printf("add ssh identity: %v", err)
+	}
 	log.Printf("uploaded generated key %s (key_id=%s)", label, keyID)
 	return nil
 }