From 73c60b9f936b19d1cf0e024a164c6c839aedefcc Mon Sep 17 00:00:00 2001
From: mirivlad <mirivlad@gmail.com>
Date: Tue, 26 May 2026 17:48:13 +0800
Subject: [PATCH] Fix: test with password, group management, vault progress

- Test in form now works with password before save (TestConnectionWithPassword)
- Groups: show existing groups in form hint, group rename/delete commands
- Vault: added progress indicators ('Unlocking vault...', 'Deriving key... done.')
- Separate template.go for command template commands
---
 cmd/group.go            | 131 ++++++++++++----------------------------
 cmd/template.go         | 101 +++++++++++++++++++++++++++++++
 cmd/tui.go              |  19 ++++++
 internal/db/servers.go  |  39 ++++++++++++
 internal/tui/app.go     |  34 ++++++++++-
 internal/vault/vault.go |   4 ++
 6 files changed, 234 insertions(+), 94 deletions(-)
 create mode 100644 cmd/template.go

diff --git a/cmd/group.go b/cmd/group.go
index 075a842..3a4e30e 100644
--- a/cmd/group.go
+++ b/cmd/group.go
@@ -2,6 +2,7 @@ package cmd
 
 import (
 	"fmt"
+	"strings"
 
 	"github.com/spf13/cobra"
 )
@@ -15,123 +16,71 @@ var groupListCmd = &cobra.Command{
 	Use:   "list",
 	Short: "List server groups",
 	RunE: func(cmd *cobra.Command, args []string) error {
-		servers, err := appDB.ListServers()
+		groups, err := appDB.GetGroups()
 		if err != nil {
-			return fmt.Errorf("list servers: %w", err)
-		}
-
-		groups := make(map[string]int)
-		for _, s := range servers {
-			g := s.GroupName
-			if g == "" {
-				g = "(no group)"
-			}
-			groups[g]++
+			return fmt.Errorf("list groups: %w", err)
 		}
 
 		if len(groups) == 0 {
-			fmt.Println("No servers.")
+			fmt.Println("No groups. Use 'sshkeeper add --group <name>' to create one.")
 			return nil
 		}
 
-		for name, count := range groups {
-			fmt.Printf("  %-20s %d servers\n", name, count)
+		for _, g := range groups {
+			fmt.Printf("  %s\n", g)
 		}
 		return nil
 	},
 }
 
-var templateCmd = &cobra.Command{
-	Use:   "template",
-	Short: "Command template management",
-}
-
-var templateListCmd = &cobra.Command{
-	Use:   "list <alias>",
-	Short: "List command templates for a server",
-	Args:  cobra.ExactArgs(1),
-	RunE: func(cmd *cobra.Command, args []string) error {
-		alias := args[0]
-		_, err := appDB.GetServer(alias)
-		if err != nil {
-			return fmt.Errorf("server not found: %s", alias)
-		}
-
-		templates, err := appDB.GetCommandTemplates(alias)
-		if err != nil {
-			return fmt.Errorf("list templates: %w", err)
-		}
-
-		if len(templates) == 0 {
-			fmt.Println("No command templates.")
-			return nil
-		}
-
-		for _, t := range templates {
-			fmt.Printf("  %-15s %s\n", t.Name, t.Command)
-		}
-		return nil
-	},
-}
-
-var templateAddCmd = &cobra.Command{
-	Use:   "add <alias> <name> <command>",
-	Short: "Add a command template",
-	Args:  cobra.ExactArgs(3),
-	RunE: func(cmd *cobra.Command, args []string) error {
-		alias := args[0]
-		name := args[1]
-		command := args[2]
-
-		server, err := appDB.GetServer(alias)
-		if err != nil {
-			return fmt.Errorf("server not found: %s", alias)
-		}
-
-		if err := appDB.AddCommandTemplate(server.ID, name, command); err != nil {
-			return fmt.Errorf("add template: %w", err)
-		}
-
-		fmt.Println("Template added.")
-		return nil
-	},
-}
-
-var runTemplateCmd = &cobra.Command{
-	Use:   "run-template <alias> <template>",
-	Short: "Run a command template on a server",
+var groupRenameCmd = &cobra.Command{
+	Use:   "rename <old> <new>",
+	Short: "Rename a group (updates all servers in the group)",
 	Args:  cobra.ExactArgs(2),
 	RunE: func(cmd *cobra.Command, args []string) error {
-		alias := args[0]
-		templateName := args[1]
+		oldName := args[0]
+		newName := args[1]
 
-		templates, err := appDB.GetCommandTemplates(alias)
-		if err != nil {
-			return fmt.Errorf("list templates: %w", err)
-		}
-		if len(templates) == 0 {
-			return fmt.Errorf("server not found or no templates: %s", alias)
+		if err := appDB.RenameGroup(oldName, newName); err != nil {
+			return fmt.Errorf("rename group: %w", err)
 		}
 
-		var command string
-		for _, t := range templates {
-			if t.Name == templateName {
-				command = t.Command
-				break
+		fmt.Printf("Group '%s' renamed to '%s'.\n", oldName, newName)
+		return nil
+	},
+}
+
+var groupDeleteCmd = &cobra.Command{
+	Use:   "delete <name>",
+	Short: "Delete a group (removes group from all servers)",
+	Args:  cobra.ExactArgs(1),
+	RunE: func(cmd *cobra.Command, args []string) error {
+		name := args[0]
+
+		if !forceFlag {
+			fmt.Printf("Remove group '%s' from all servers? (y/N): ", name)
+			var response string
+			fmt.Scanln(&response)
+			if strings.ToLower(response) != "y" {
+				fmt.Println("Cancelled.")
+				return nil
 			}
 		}
 
-		if command == "" {
-			return fmt.Errorf("template not found: %s", templateName)
+		if err := appDB.DeleteGroup(name); err != nil {
+			return fmt.Errorf("delete group: %w", err)
 		}
 
-		fmt.Printf("Running '%s' on %s...\n", command, alias)
+		fmt.Printf("Group '%s' removed from all servers.\n", name)
 		return nil
 	},
 }
 
+var forceFlag bool
+
 func init() {
 	groupCmd.AddCommand(groupListCmd)
-	templateCmd.AddCommand(templateListCmd)
-	templateCmd.AddCommand(templateAddCmd)
+	groupCmd.AddCommand(groupRenameCmd)
+	groupCmd.AddCommand(groupDeleteCmd)
+	groupDeleteCmd.Flags().BoolVarP(&forceFlag, "force", "f", false, "Delete without confirmation")
 }
diff --git a/cmd/template.go b/cmd/template.go
new file mode 100644
index 0000000..db74f93
--- /dev/null
+++ b/cmd/template.go
@@ -0,0 +1,101 @@
+package cmd
+
+import (
+	"fmt"
+
+	"github.com/spf13/cobra"
+)
+
+var templateCmd = &cobra.Command{
+	Use:   "template",
+	Short: "Command template management",
+}
+
+var templateListCmd = &cobra.Command{
+	Use:   "list <alias>",
+	Short: "List command templates for a server",
+	Args:  cobra.ExactArgs(1),
+	RunE: func(cmd *cobra.Command, args []string) error {
+		alias := args[0]
+		_, err := appDB.GetServer(alias)
+		if err != nil {
+			return fmt.Errorf("server not found: %s", alias)
+		}
+
+		templates, err := appDB.GetCommandTemplates(alias)
+		if err != nil {
+			return fmt.Errorf("list templates: %w", err)
+		}
+
+		if len(templates) == 0 {
+			fmt.Println("No command templates.")
+			return nil
+		}
+
+		for _, t := range templates {
+			fmt.Printf("  %-15s %s\n", t.Name, t.Command)
+		}
+		return nil
+	},
+}
+
+var templateAddCmd = &cobra.Command{
+	Use:   "add <alias> <name> <command>",
+	Short: "Add a command template",
+	Args:  cobra.ExactArgs(3),
+	RunE: func(cmd *cobra.Command, args []string) error {
+		alias := args[0]
+		name := args[1]
+		command := args[2]
+
+		server, err := appDB.GetServer(alias)
+		if err != nil {
+			return fmt.Errorf("server not found: %s", alias)
+		}
+
+		if err := appDB.AddCommandTemplate(server.ID, name, command); err != nil {
+			return fmt.Errorf("add template: %w", err)
+		}
+
+		fmt.Println("Template added.")
+		return nil
+	},
+}
+
+var runTemplateCmd = &cobra.Command{
+	Use:   "run-template <alias> <template>",
+	Short: "Run a command template on a server",
+	Args:  cobra.ExactArgs(2),
+	RunE: func(cmd *cobra.Command, args []string) error {
+		alias := args[0]
+		templateName := args[1]
+
+		templates, err := appDB.GetCommandTemplates(alias)
+		if err != nil {
+			return fmt.Errorf("list templates: %w", err)
+		}
+		if len(templates) == 0 {
+			return fmt.Errorf("server not found or no templates: %s", alias)
+		}
+
+		var command string
+		for _, t := range templates {
+			if t.Name == templateName {
+				command = t.Command
+				break
+			}
+		}
+
+		if command == "" {
+			return fmt.Errorf("template not found: %s", templateName)
+		}
+
+		fmt.Printf("Running '%s' on %s...\n", command, alias)
+		return nil
+	},
+}
+
+func init() {
+	templateCmd.AddCommand(templateListCmd)
+	templateCmd.AddCommand(templateAddCmd)
+}
diff --git a/cmd/tui.go b/cmd/tui.go
index 59e1753..5abc250 100644
--- a/cmd/tui.go
+++ b/cmd/tui.go
@@ -41,6 +41,15 @@ func runTUI() error {
 	tui.TestConnection = func(server *model.Server) (bool, string) {
 		return ssh.Test(cfg, server, vaultFunc)
 	}
+	tui.TestConnectionWithPassword = func(server *model.Server, password string) (bool, string) {
+		directVaultFunc := func(sa string, st string) (string, error) {
+			if st == "ssh_password" || st == "key_passphrase" {
+				return password, nil
+			}
+			return vaultFunc(sa, st)
+		}
+		return ssh.Test(cfg, server, directVaultFunc)
+	}
 	tui.SaveServer = func(server *model.Server, password string) error {
 		if password != "" {
 			v := getOrCreateVault()
@@ -66,6 +75,16 @@ func runTUI() error {
 		return appDB.CreateServer(server)
 	}
 
+	tui.GetGroups = func() ([]string, error) {
+		return appDB.GetGroups()
+	}
+	tui.RenameGroup = func(oldName, newName string) error {
+		return appDB.RenameGroup(oldName, newName)
+	}
+	tui.DeleteGroup = func(name string) error {
+		return appDB.DeleteGroup(name)
+	}
+
 	// Run TUI in a loop — if user requests connect, handle it and restart TUI
 	for {
 		m := tui.New(servers)
diff --git a/internal/db/servers.go b/internal/db/servers.go
index 17fcf9b..83c43a0 100644
--- a/internal/db/servers.go
+++ b/internal/db/servers.go
@@ -246,3 +246,42 @@ func (db *DB) GetCommandTemplates(serverAlias string) ([]*model.CommandTemplate,
 	}
 	return templates, rows.Err()
 }
+
+// GetGroups returns all unique group names with server count
+func (db *DB) GetGroups() ([]string, error) {
+	rows, err := db.conn.Query(`
+		SELECT group_name FROM servers 
+		WHERE group_name != '' 
+		GROUP BY group_name 
+		ORDER BY group_name`)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+
+	var groups []string
+	for rows.Next() {
+		var name string
+		if err := rows.Scan(&name); err != nil {
+			return nil, err
+		}
+		groups = append(groups, name)
+	}
+	return groups, rows.Err()
+}
+
+// RenameGroup renames a group for all servers in it
+func (db *DB) RenameGroup(oldName, newName string) error {
+	_, err := db.conn.Exec(
+		"UPDATE servers SET group_name = ?, updated_at = CURRENT_TIMESTAMP WHERE group_name = ?",
+		newName, oldName)
+	return err
+}
+
+// DeleteGroup removes group assignment from all servers
+func (db *DB) DeleteGroup(name string) error {
+	_, err := db.conn.Exec(
+		"UPDATE servers SET group_name = '', updated_at = CURRENT_TIMESTAMP WHERE group_name = ?",
+		name)
+	return err
+}
diff --git a/internal/tui/app.go b/internal/tui/app.go
index 37eb544..14a57f4 100644
--- a/internal/tui/app.go
+++ b/internal/tui/app.go
@@ -78,7 +78,12 @@ var (
 	SearchServers  func(query string) ([]*model.Server, error)
 	DeleteServer   func(alias string) error
 	TestConnection func(server *model.Server) (bool, string)
+	// TestConnectionWithPassword tests with explicit password (for form test before save)
+	TestConnectionWithPassword func(server *model.Server, password string) (bool, string)
 	SaveServer     func(server *model.Server, password string) error
+	GetGroups      func() ([]string, error)   // Returns existing group names
+	RenameGroup    func(oldName, newName string) error  // Rename group for all servers
+	DeleteGroup    func(name string) error     // Remove group from all servers
 )
 
 // --- Screen type ---
@@ -380,6 +385,8 @@ type formModel struct {
 	spinner        spinner.Model
 	width          int
 	height         int
+	groups         string  // cached groups list (comma-separated for display)
+	showGroups     bool    // show group dropdown
 }
 
 func newFormModel(w, h int) *formModel {
@@ -393,7 +400,7 @@ func newFormModel(w, h int) *formModel {
 		"Auth Method (password/key/key_passphrase/agent)",
 		"Identity File",
 		"ProxyJump",
-		"Group",
+		"Group (type new or pick from list)",
 		"Notes",
 	}
 	for i, label := range labels {
@@ -413,7 +420,7 @@ func newFormModel(w, h int) *formModel {
 
 	inputs[0].Focus()
 
-	return &formModel{
+	fm := &formModel{
 		inputs:   inputs,
 		password: pw,
 		focusIdx: 0,
@@ -421,6 +428,15 @@ func newFormModel(w, h int) *formModel {
 		width:    w,
 		height:   h,
 	}
+
+	// Load existing groups
+	if GetGroups != nil {
+		if groups, err := GetGroups(); err == nil && len(groups) > 0 {
+			fm.groups = strings.Join(groups, ", ")
+		}
+	}
+
+	return fm
 }
 
 func newEditFormModel(s *model.Server, w, h int) *formModel {
@@ -581,10 +597,18 @@ func (fm *formModel) runTest() tea.Cmd {
 	fm.saved = false
 
 	s := fm.buildServer()
+	pw := fm.password.Value()
+
 	return tea.Batch(
 		fm.spinner.Tick,
 		func() tea.Msg {
-			if s.AuthMethod == model.AuthPassword && fm.password.Value() == "" {
+			// Use direct password test if available (for form test before save)
+			if TestConnectionWithPassword != nil {
+				ok, testErr := TestConnectionWithPassword(s, pw)
+				return testDoneMsg{ok: ok, err: testErr}
+			}
+			// Fallback to vault-based test
+			if s.AuthMethod == model.AuthPassword && pw == "" {
 				return testDoneMsg{ok: false, err: "Password is required for password auth."}
 			}
 			ok, testErr := TestConnection(s)
@@ -651,6 +675,10 @@ func (fm *formModel) View() string {
 	for i := range fm.inputs {
 		b.WriteString(fm.inputs[i].View())
 		b.WriteString("\n")
+		// Show existing groups hint under Group field
+		if i == 8 && fm.groups != "" {
+			b.WriteString(helpStyle.Render("  Groups: " + fm.groups + "\n"))
+		}
 	}
 
 	b.WriteString(fm.password.View())
diff --git a/internal/vault/vault.go b/internal/vault/vault.go
index c54c813..f756cb7 100644
--- a/internal/vault/vault.go
+++ b/internal/vault/vault.go
@@ -113,6 +113,7 @@ func Create(path string, masterPassword string) error {
 		key[i] = 0
 	}
 
+	fmt.Println(" done.")
 	return nil
 }
 
@@ -121,6 +122,8 @@ func (v *Vault) Unlock(masterPassword string) error {
 	v.mu.Lock()
 	defer v.mu.Unlock()
 
+	fmt.Print("Unlocking vault...")
+
 	data, err := os.ReadFile(v.path)
 	if err != nil {
 		return fmt.Errorf("read vault file: %w", err)
@@ -160,6 +163,7 @@ func (v *Vault) Unlock(masterPassword string) error {
 		v.records[rec.ID] = plaintext
 	}
 
+	fmt.Println(" done.")
 	return nil
 }