From f68d7c2c947203855f02f3412310d5469c2f1240 Mon Sep 17 00:00:00 2001
From: juancwu <46619361+juancwu@users.noreply.github.com>
Date: Sun, 11 Jan 2026 18:09:44 -0500
Subject: [PATCH] update key in db

---
 agent.go |  2 +-
 cmd.go   | 15 +++++++++++++++
 db.go    | 28 ++++++++++++++++++++++++++++
 main.go  | 14 ++++++++++++++
 4 files changed, 58 insertions(+), 1 deletion(-)

diff --git a/agent.go b/agent.go
index e716493..ceee523 100644
--- a/agent.go
+++ b/agent.go
@@ -23,7 +23,7 @@ func startEphemeralAgent(pemData []byte, target string) (string, func(), error)
 
 	if err != nil {
 		if _, ok := err.(*ssh.PassphraseMissingError); ok {
-			fmt.Printf("\033[1;32m? Gosh:\033[0m Key for \033[1m%s\033[0m is encrypted. Enter passphrase: ", target)
+			fmt.Printf("\033[1;32mGosh:\033[0m Key for \033[1m%s\033[0m is encrypted. Enter passphrase: ", target)
 			pass, readErr := term.ReadPassword(int(syscall.Stdin))
 			fmt.Println()
 			if readErr != nil {
diff --git a/cmd.go b/cmd.go
index 526e9c4..d004790 100644
--- a/cmd.go
+++ b/cmd.go
@@ -41,3 +41,18 @@ func handleListKey(dbPath string) error {
 	w.Flush()
 	return nil
 }
+
+func handleUpdateKey(dbPath, userPattern, hostPattern, keyPath string) error {
+	db, err := initDB(dbPath)
+	if err != nil {
+		return err
+	}
+	defer db.Close()
+
+	err = updateKey(db, userPattern, hostPattern, keyPath)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
diff --git a/db.go b/db.go
index b29bf04..7f7f53b 100644
--- a/db.go
+++ b/db.go
@@ -93,3 +93,31 @@ func listkeys(db *sql.DB) ([]KeyRecord, error) {
 
 	return keys, nil
 }
+
+func updateKey(db *sql.DB, userPattern, hostPattern, keyPath string) error {
+	pemData, err := os.ReadFile(keyPath)
+	if err != nil {
+		return fmt.Errorf("failed to read key: %w", err)
+	}
+
+	res, err := db.Exec(
+		"UPDATE keys SET encrypted_pem=?, comment=? WHERE user_pattern = ? AND host_pattern = ?;",
+		pemData, "Updated from "+keyPath, userPattern, hostPattern,
+	)
+	if err != nil {
+		return fmt.Errorf("failed to update key: %w", err)
+	}
+
+	rows, err := res.RowsAffected()
+	if err == nil {
+		if rows == 0 {
+			fmt.Printf("No key found with user '%s' and host '%s'.\n", userPattern, hostPattern)
+		} else {
+			fmt.Printf("Key for %s@%s updated successfully.\n", userPattern, hostPattern)
+		}
+	} else {
+		fmt.Println("Warning: could not verify update result.", err)
+	}
+
+	return nil
+}
diff --git a/main.go b/main.go
index 8fb79dd..02ab17c 100644
--- a/main.go
+++ b/main.go
@@ -12,6 +12,7 @@ func main() {
 		fmt.Fprintf(os.Stderr, "\nManagement Commands:\n")
 		fmt.Fprintf(os.Stderr, "  %s [flags] list-keys\n", os.Args[0])
 		fmt.Fprintf(os.Stderr, "  %s [flags] add-key <user_pattern> <host_pattern> <path_to_key>\n", os.Args[0])
+		fmt.Fprintf(os.Stderr, "  %s [flags] update-key <user_pattern> <host_pattern> <path_to_key>\n", os.Args[0])
 		fmt.Fprintf(os.Stderr, "\nFlags:\n")
 		flag.PrintDefaults()
 	}
@@ -51,6 +52,19 @@ func main() {
 			os.Exit(1)
 		}
 		os.Exit(0)
+
+	case "update-key":
+		if argc != 4 {
+			fmt.Println("Usage: gosh update-key <user_pattern> <host_pattern> <path_to_key>")
+			os.Exit(1)
+		}
+
+		err := handleUpdateKey(*dbPath, args[1], args[2], args[3])
+		if err != nil {
+			fmt.Println("Error:", err)
+			os.Exit(1)
+		}
+		os.Exit(0)
 	}
 
 	startSSH(*dbPath, args)