feat: Cleanup stale keys on VM destroy

Author: theognis1002

cmd/destroy.go

@@ -113,23 +113,19 @@ Examples:
 
 		fmt.Println()
 
-		// Check if we should destroy the VM
 		// Only destroy if: (1) VM was provisioned by this target AND (2) no other apps on the server
 		shouldDestroyVM := false
 		var otherApps []state.DeployedApp
 
 		if provisionedID != "" && target.Provider != "" && target.Provider != "byos" {
-			// Check if this target provisioned the VM
 			targetProvisionedVM := false
 			if providerCfg != nil {
 				targetProvisionedVM = providerCfg.IsProvisioned()
 			}
 
-			// Check if other apps exist on this server
 			if target.ServerIP != "" {
 				serverState, err := state.GetServerState(target.ServerIP)
 				if err == nil {
-					// Get apps other than the current one
 					for _, app := range serverState.DeployedApps {
 						if app.TargetName != destroyTargetFlag {
 							otherApps = append(otherApps, app)
@@ -138,12 +134,9 @@ Examples:
 				}
 			}
 
-			// Decide whether to destroy VM
 			if !targetProvisionedVM {
-				// This target didn't provision the VM (joined existing server)
 				shouldDestroyVM = false
 			} else if len(otherApps) > 0 {
-				// Other apps exist on this server
 				shouldDestroyVM = false
 				fmt.Printf("%s %s\n", destroyWarningStyle.Render("⚠"), destroyWarningStyle.Render("VM will NOT be destroyed - other apps are deployed to this server:"))
 				for _, app := range otherApps {
@@ -153,7 +146,6 @@ Examples:
 				fmt.Printf("%s %s\n", destroyMutedStyle.Render("ℹ"), destroyMutedStyle.Render("Only removing this app's configuration and local state"))
 				fmt.Println()
 			} else {
-				// This target provisioned the VM and no other apps exist
 				shouldDestroyVM = true
 			}
 		}
@@ -224,14 +216,12 @@ Examples:
 			}
 		}
 
-		// Unregister app from server state if applicable
 		if target.ServerIP != "" {
 			if err := state.UnregisterApp(target.ServerIP, destroyTargetFlag); err != nil {
 				fmt.Printf("%s %s\n", destroyWarningStyle.Render("⚠"), destroyMutedStyle.Render(fmt.Sprintf("Failed to unregister app from server: %v", err)))
 			} else {
 				fmt.Printf("%s %s\n", destroySuccessStyle.Render("✓"), destroyMutedStyle.Render("Unregistered app from server state"))
 
-				// Check if other apps remain on this server
 				serverState, err := state.GetServerState(target.ServerIP)
 				if err == nil {
 					if len(serverState.DeployedApps) > 0 {
@@ -241,17 +231,14 @@ Examples:
 
 				// Clean up unused runtimes if keeping the VM (only makes sense for multi-app servers)
 				if !shouldDestroyVM && len(otherApps) > 0 {
-					// Get provider config to connect via SSH
 					providerCfg, err := target.GetAnyProviderConfig()
 					if err == nil && providerCfg.GetIP() != "" {
 						fmt.Printf("%s %s\n", destroyMutedStyle.Render("→"), destroyMutedStyle.Render("Analyzing runtime dependencies..."))
 
-						// Connect to server via SSH
 						sshExecutor := sshpkg.NewExecutor(providerCfg.GetIP(), "22", providerCfg.GetUsername(), providerCfg.GetSSHKey())
 						if err := sshExecutor.Connect(3, 10*time.Second); err == nil {
 							defer sshExecutor.Disconnect()
 
-							// Cleanup unused runtimes
 							if err := runtime.CleanupUnusedRuntimes(sshExecutor, target.ServerIP, destroyTargetFlag); err != nil {
 								fmt.Printf("%s %s\n", destroyWarningStyle.Render("⚠"), destroyMutedStyle.Render(fmt.Sprintf("Runtime cleanup warning: %v", err)))
 								fmt.Printf("%s %s\n", destroyMutedStyle.Render("  "), destroyMutedStyle.Render("You may manually run: apt-get autoremove -y"))
@@ -278,11 +265,21 @@ Examples:
 		}
 		fmt.Printf("%s %s\n", destroySuccessStyle.Render("✓"), destroyMutedStyle.Render("Removed target from config"))
 
+		fmt.Printf("%s %s\n", destroyMutedStyle.Render("→"), destroyMutedStyle.Render("Checking for unused SSH keys..."))
+		keysDeleted, err := sshpkg.CleanupUnusedKeys(cfg.Targets)
+		if err != nil {
+			fmt.Printf("%s %s\n", destroyWarningStyle.Render("⚠"), destroyMutedStyle.Render(fmt.Sprintf("SSH key cleanup warning: %v", err)))
+		} else if keysDeleted > 0 {
+			fmt.Printf("%s %s\n", destroySuccessStyle.Render("✓"), destroyMutedStyle.Render(fmt.Sprintf("Cleaned up %d unused SSH key(s)", keysDeleted)))
+		} else {
+			fmt.Printf("%s %s\n", destroyMutedStyle.Render("ℹ"), destroyMutedStyle.Render("No unused SSH keys found"))
+		}
+
 		fmt.Println()
 
 		successBox := lipgloss.NewStyle().
 			Border(lipgloss.RoundedBorder()).
-			BorderForeground(lipgloss.Color("196")). // Red border for destruction
+			BorderForeground(lipgloss.Color("196")).
 			Padding(0, 1).
 			Render(
 				lipgloss.JoinVertical(

pkg/ssh/keygen.go

@@ -24,19 +24,16 @@ type KeyPair struct {
 
 // GenerateKeyPair generates a new Ed25519 SSH key pair
 func GenerateKeyPair(keyName string) (*KeyPair, error) {
-	// Generate Ed25519 key pair
 	publicKey, privateKey, err := ed25519.GenerateKey(rand.Reader)
 	if err != nil {
 		return nil, fmt.Errorf("failed to generate key pair: %w", err)
 	}
 
-	// Create SSH public key
 	sshPublicKey, err := ssh.NewPublicKey(publicKey)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create SSH public key: %w", err)
 	}
 
-	// Convert private key to PEM format
 	privateKeyBytes, err := x509.MarshalPKCS8PrivateKey(privateKey)
 	if err != nil {
 		return nil, fmt.Errorf("failed to marshal private key: %w", err)
@@ -47,29 +44,23 @@ func GenerateKeyPair(keyName string) (*KeyPair, error) {
 		Bytes: privateKeyBytes,
 	})
 
-	// Format SSH public key
 	sshPublicKeyBytes := ssh.MarshalAuthorizedKey(sshPublicKey)
 	sshPublicKeyString := strings.TrimSpace(string(sshPublicKeyBytes))
 
-	// Get fingerprint
 	fingerprint := ssh.FingerprintSHA256(sshPublicKey)
 
-	// Get lightfold keys directory
 	keysDir, err := GetKeysDirectory()
 	if err != nil {
 		return nil, fmt.Errorf("failed to get keys directory: %w", err)
 	}
 
-	// Create key file paths
 	privateKeyPath := filepath.Join(keysDir, keyName)
 	publicKeyPath := privateKeyPath + ".pub"
 
-	// Write private key
 	if err := os.WriteFile(privateKeyPath, privateKeyPEM, config.PermPrivateKey); err != nil {
 		return nil, fmt.Errorf("failed to write private key: %w", err)
 	}
 
-	// Write public key
 	if err := os.WriteFile(publicKeyPath, sshPublicKeyBytes, config.PermPublicKey); err != nil {
 		return nil, fmt.Errorf("failed to write public key: %w", err)
 	}
@@ -91,7 +82,6 @@ func GetKeysDirectory() (string, error) {
 
 	keysDir := filepath.Join(homeDir, config.LocalConfigDir, config.LocalKeysDir)
 
-	// Create directory if it doesn't exist
 	if err := os.MkdirAll(keysDir, 0700); err != nil {
 		return "", fmt.Errorf("failed to create keys directory: %w", err)
 	}
@@ -111,24 +101,19 @@ func LoadPublicKey(publicKeyPath string) (string, error) {
 
 // ValidateSSHKey validates that an SSH key file exists and is valid
 func ValidateSSHKey(keyPath string) error {
-	// Check if file exists
 	if _, err := os.Stat(keyPath); os.IsNotExist(err) {
 		return fmt.Errorf("SSH key file does not exist: %s", keyPath)
 	}
 
-	// Try to read and parse the key
 	data, err := os.ReadFile(keyPath)
 	if err != nil {
 		return fmt.Errorf("failed to read SSH key file: %w", err)
 	}
 
-	// Try to parse as private key first
 	block, _ := pem.Decode(data)
 	if block != nil {
-		// It's a PEM-encoded private key
 		_, err := x509.ParsePKCS8PrivateKey(block.Bytes)
 		if err != nil {
-			// Try parsing as RSA private key
 			_, err = x509.ParsePKCS1PrivateKey(block.Bytes)
 			if err != nil {
 				return fmt.Errorf("invalid private key format: %w", err)
@@ -137,7 +122,6 @@ func ValidateSSHKey(keyPath string) error {
 		return nil
 	}
 
-	// Try to parse as SSH public key
 	_, _, _, _, err = ssh.ParseAuthorizedKey(data)
 	if err != nil {
 		return fmt.Errorf("invalid SSH key format: %w", err)
@@ -148,7 +132,6 @@ func ValidateSSHKey(keyPath string) error {
 
 // GetKeyName generates a key name for a project
 func GetKeyName(projectName string) string {
-	// Sanitize project name for file system
 	keyName := strings.ReplaceAll(projectName, " ", "_")
 	keyName = strings.ReplaceAll(keyName, "/", "_")
 	keyName = strings.ReplaceAll(keyName, "\\", "_")
@@ -167,7 +150,6 @@ func KeyExists(keyName string) (bool, error) {
 	privateKeyPath := filepath.Join(keysDir, keyName)
 	publicKeyPath := privateKeyPath + ".pub"
 
-	// Check if both files exist
 	if _, err := os.Stat(privateKeyPath); os.IsNotExist(err) {
 		return false, nil
 	}
@@ -177,3 +159,83 @@ func KeyExists(keyName string) (bool, error) {
 
 	return true, nil
 }
+
+// DeleteKeyPair deletes an SSH key pair
+func DeleteKeyPair(keyPath string) error {
+	if err := os.Remove(keyPath); err != nil && !os.IsNotExist(err) {
+		return fmt.Errorf("failed to delete private key: %w", err)
+	}
+
+	publicKeyPath := keyPath + ".pub"
+	if err := os.Remove(publicKeyPath); err != nil && !os.IsNotExist(err) {
+		return fmt.Errorf("failed to delete public key: %w", err)
+	}
+
+	return nil
+}
+
+// GetSSHKeyFromPath extracts the base key name from a full path
+// Example: "/home/user/.lightfold/keys/lightfold_my-project_ed25519" -> "lightfold_my-project_ed25519"
+func GetSSHKeyFromPath(keyPath string) string {
+	return filepath.Base(keyPath)
+}
+
+// CleanupUnusedKeys removes SSH keys that are not used by any targets.
+// Returns the number of keys deleted.
+func CleanupUnusedKeys(allTargets map[string]config.TargetConfig) (int, error) {
+	keysDir, err := GetKeysDirectory()
+	if err != nil {
+		return 0, fmt.Errorf("failed to get keys directory: %w", err)
+	}
+
+	keysInUse := make(map[string]bool)
+
+	for _, target := range allTargets {
+		providerCfg, err := target.GetSSHProviderConfig()
+		if err != nil || providerCfg == nil {
+			continue
+		}
+
+		sshKey := providerCfg.GetSSHKey()
+		if sshKey != "" {
+			keyName := filepath.Base(sshKey)
+			keysInUse[keyName] = true
+		}
+	}
+
+	entries, err := os.ReadDir(keysDir)
+	if err != nil {
+		return 0, fmt.Errorf("failed to read keys directory: %w", err)
+	}
+
+	var keysToDelete []string
+
+	for _, entry := range entries {
+		if entry.IsDir() {
+			continue
+		}
+
+		filename := entry.Name()
+
+		// Skip public keys - we'll delete them with their private keys
+		if strings.HasSuffix(filename, ".pub") {
+			continue
+		}
+
+		if !strings.HasPrefix(filename, "lightfold_") || !strings.HasSuffix(filename, "_ed25519") {
+			continue
+		}
+
+		if !keysInUse[filename] {
+			keysToDelete = append(keysToDelete, filepath.Join(keysDir, filename))
+		}
+	}
+
+	for _, keyPath := range keysToDelete {
+		if err := DeleteKeyPair(keyPath); err != nil {
+			return len(keysToDelete), fmt.Errorf("failed to delete key %s: %w", filepath.Base(keyPath), err)
+		}
+	}
+
+	return len(keysToDelete), nil
+}