From 573ccff5cb1d394d2004e504b5034aff3e3a5e18 Mon Sep 17 00:00:00 2001 From: Hailong Yang Date: Sun, 8 Sep 2024 16:16:13 -0400 Subject: [PATCH 1/2] fix: only add key when not exist in db. --- cmd/config/config.go | 79 ++++++++++++++++++++++++++++++++++++-------- cmd/ssh/functions.go | 20 +++++++++++ 2 files changed, 86 insertions(+), 13 deletions(-) diff --git a/cmd/config/config.go b/cmd/config/config.go index 1f2344c..a201a33 100644 --- a/cmd/config/config.go +++ b/cmd/config/config.go @@ -2,7 +2,7 @@ package config import ( "fmt" - "os" + "strings" "github.com/runpod/runpodctl/api" "github.com/runpod/runpodctl/cmd/ssh" @@ -21,27 +21,80 @@ var ConfigCmd = &cobra.Command{ Use: "config", Short: "Manage CLI configuration", Long: "RunPod CLI Config Settings", - Run: func(c *cobra.Command, args []string) { - if err := viper.WriteConfig(); err != nil { - fmt.Fprintf(os.Stderr, "Error saving config: %v\n", err) - return + RunE: func(c *cobra.Command, args []string) error { + if err := saveConfig(); err != nil { + return fmt.Errorf("error saving config: %w", err) } - fmt.Println("Configuration saved to file:", viper.ConfigFileUsed()) - publicKey, err := ssh.GenerateSSHKeyPair("RunPod-Key-Go") + publicKey, err := getOrCreateSSHKey() if err != nil { - fmt.Fprintf(os.Stderr, "Failed to generate SSH key: %v\n", err) - return + return fmt.Errorf("failed to get or create local SSH key: %w", err) } - if err := api.AddPublicSSHKey(publicKey); err != nil { - fmt.Fprintf(os.Stderr, "Failed to add the SSH key: %v\n", err) - return + if err := updateSSHKeyInCloud(publicKey); err != nil { + return fmt.Errorf("failed to update SSH key in the cloud: %w", err) } - fmt.Println("SSH key added successfully.") + + return nil }, } +// saveConfig saves the CLI configuration to a file +func saveConfig() error { + if err := viper.WriteConfig(); err != nil { + return err + } + fmt.Println("Configuration saved to file:", viper.ConfigFileUsed()) + return nil +} + +// Checks for an existing local SSH key and generates a new one if not found +func getOrCreateSSHKey() ([]byte, error) { + publicKey, err := ssh.GetLocalSSHKey() + if err != nil { + return nil, fmt.Errorf("error checking for local SSH key: %w", err) + } + + if publicKey == nil { + fmt.Println("No existing local SSH key found, generating a new one.") + publicKey, err = ssh.GenerateSSHKeyPair("RunPod-Key-Go") + if err != nil { + return nil, fmt.Errorf("failed to generate SSH key: %w", err) + } + fmt.Println("New SSH key pair generated.") + } else { + fmt.Println("Existing local SSH key found.") + } + + return publicKey, nil +} + +// Checks if the SSH key exists in the cloud and adds it if necessary +func updateSSHKeyInCloud(publicKey []byte) error { + _, cloudKeys, err := api.GetPublicSSHKeys() + if err != nil { + return fmt.Errorf("failed to get SSH key in the cloud: %w", err) + } + + // Trim any whitespace from the publicKey + publicKeyStr := strings.TrimSpace(string(publicKey)) + + // Check if the publicKey already exists in the cloud + for _, cloudKey := range cloudKeys { + if strings.TrimSpace(cloudKey.Key) == publicKeyStr { + fmt.Println("SSH key already exists in the cloud. No action needed.") + return nil + } + } + + if err := api.AddPublicSSHKey(publicKey); err != nil { + return fmt.Errorf("failed to add the SSH key: %w", err) + } + + fmt.Println("SSH key added successfully to the cloud.") + return nil +} + func init() { ConfigCmd.Flags().StringVar(&apiKey, "apiKey", "", "RunPod API key") viper.BindPFlag("apiKey", ConfigCmd.Flags().Lookup("apiKey")) //nolint diff --git a/cmd/ssh/functions.go b/cmd/ssh/functions.go index d613833..0185dfe 100644 --- a/cmd/ssh/functions.go +++ b/cmd/ssh/functions.go @@ -63,3 +63,23 @@ func GenerateSSHKeyPair(keyName string) ([]byte, error) { fmt.Printf("SSH key pair generated: %s (private), %s (public)\n", privateKeyPath, publicKeyPath) return publicKeyBytes, nil } + +func GetLocalSSHKey() ([]byte, error) { + usr, err := os.UserHomeDir() + if err != nil { + return nil, fmt.Errorf("failed to get current user: %w", err) + } + + keyPath := filepath.Join(usr, ".ssh", "RunPod-Key-Go.pub") + + if _, err := os.Stat(keyPath); os.IsNotExist(err) { + return nil, nil // No existing key found + } + + publicKey, err := os.ReadFile(keyPath) + if err != nil { + return nil, fmt.Errorf("failed to read existing public key: %w", err) + } + + return publicKey, nil +} From a9325873d5469d38e28afc18e14bf7f2c3d837e4 Mon Sep 17 00:00:00 2001 From: Hailong Yang Date: Sun, 8 Sep 2024 20:33:43 -0400 Subject: [PATCH 2/2] fix: ssh key path and comparsion --- cmd/config/config.go | 22 ++++++++++++++-------- cmd/ssh/functions.go | 19 ++++++++++++------- 2 files changed, 26 insertions(+), 15 deletions(-) diff --git a/cmd/config/config.go b/cmd/config/config.go index a201a33..3a3b946 100644 --- a/cmd/config/config.go +++ b/cmd/config/config.go @@ -2,13 +2,13 @@ package config import ( "fmt" - "strings" "github.com/runpod/runpodctl/api" "github.com/runpod/runpodctl/cmd/ssh" "github.com/spf13/cobra" "github.com/spf13/viper" + sshcrypto "golang.org/x/crypto/ssh" ) var ( @@ -31,7 +31,7 @@ var ConfigCmd = &cobra.Command{ return fmt.Errorf("failed to get or create local SSH key: %w", err) } - if err := updateSSHKeyInCloud(publicKey); err != nil { + if err := ensureSSHKeyInCloud(publicKey); err != nil { return fmt.Errorf("failed to update SSH key in the cloud: %w", err) } @@ -69,24 +69,30 @@ func getOrCreateSSHKey() ([]byte, error) { return publicKey, nil } -// Checks if the SSH key exists in the cloud and adds it if necessary -func updateSSHKeyInCloud(publicKey []byte) error { +// ensureSSHKeyInCloud checks if the SSH key exists in the cloud and adds it if necessary +func ensureSSHKeyInCloud(publicKey []byte) error { _, cloudKeys, err := api.GetPublicSSHKeys() if err != nil { - return fmt.Errorf("failed to get SSH key in the cloud: %w", err) + return fmt.Errorf("failed to get SSH keys from the cloud: %w", err) } - // Trim any whitespace from the publicKey - publicKeyStr := strings.TrimSpace(string(publicKey)) + // Parse the local public key + localPubKey, _, _, _, err := sshcrypto.ParseAuthorizedKey(publicKey) + if err != nil { + return fmt.Errorf("failed to parse local public key: %w", err) + } + + localFingerprint := sshcrypto.FingerprintSHA256(localPubKey) // Check if the publicKey already exists in the cloud for _, cloudKey := range cloudKeys { - if strings.TrimSpace(cloudKey.Key) == publicKeyStr { + if cloudKey.Fingerprint == localFingerprint { fmt.Println("SSH key already exists in the cloud. No action needed.") return nil } } + // If the key doesn't exist, add it if err := api.AddPublicSSHKey(publicKey); err != nil { return fmt.Errorf("failed to add the SSH key: %w", err) } diff --git a/cmd/ssh/functions.go b/cmd/ssh/functions.go index 0185dfe..3e7625c 100644 --- a/cmd/ssh/functions.go +++ b/cmd/ssh/functions.go @@ -65,21 +65,26 @@ func GenerateSSHKeyPair(keyName string) ([]byte, error) { } func GetLocalSSHKey() ([]byte, error) { - usr, err := os.UserHomeDir() + homeDir, err := os.UserHomeDir() if err != nil { - return nil, fmt.Errorf("failed to get current user: %w", err) + return nil, fmt.Errorf("failed to get user home directory: %w", err) } - keyPath := filepath.Join(usr, ".ssh", "RunPod-Key-Go.pub") - - if _, err := os.Stat(keyPath); os.IsNotExist(err) { - return nil, nil // No existing key found - } + keyPath := filepath.Join(homeDir, ".runpod", "ssh", "RunPod-Key-Go.pub") publicKey, err := os.ReadFile(keyPath) if err != nil { + if os.IsNotExist(err) { + return nil, nil // No existing key found + } return nil, fmt.Errorf("failed to read existing public key: %w", err) } + // Validate the key format + _, _, _, _, err = ssh.ParseAuthorizedKey(publicKey) + if err != nil { + return nil, fmt.Errorf("invalid public key format: %w", err) + } + return publicKey, nil }