@@ -2,6 +2,9 @@ package vms
22
33import (
44 "context"
5+ "crypto/ed25519"
6+ "crypto/rand"
7+ "encoding/pem"
58 "fmt"
69 "os"
710 "os/exec"
@@ -10,6 +13,7 @@ import (
1013
1114 "github.com/pterm/pterm"
1215 "github.com/urfave/cli/v2"
16+ gossh "golang.org/x/crypto/ssh"
1317
1418 "github.com/NodeOps-app/createos-cli/internal/api"
1519 "github.com/NodeOps-app/createos-cli/internal/terminal"
@@ -43,20 +47,65 @@ func findPublicKeys() map[string]string {
4347 return keys
4448}
4549
46- // selectPublicKey prompts the user to pick a local SSH public key, or skip.
47- // Returns the selected key content, or "" if skipped.
48- func selectPublicKey () (string , error ) {
49- keys := findPublicKeys ()
50- if len (keys ) == 0 {
51- return "" , nil
50+ // generateTempSSHKeypair creates a temporary ed25519 keypair, writes the private key
51+ // to a temp file, and returns the public key string, the private key path, and a
52+ // cleanup func that removes the temp file.
53+ func generateTempSSHKeypair () (publicKey string , privateKeyPath string , cleanup func (), err error ) {
54+ pub , priv , err := ed25519 .GenerateKey (rand .Reader )
55+ if err != nil {
56+ return "" , "" , nil , fmt .Errorf ("could not generate keypair: %w" , err )
57+ }
58+
59+ // Marshal public key to authorized_keys format.
60+ sshPub , err := gossh .NewPublicKey (pub )
61+ if err != nil {
62+ return "" , "" , nil , fmt .Errorf ("could not encode public key: %w" , err )
63+ }
64+ pubKeyStr := strings .TrimSpace (string (gossh .MarshalAuthorizedKey (sshPub )))
65+
66+ // Marshal private key to OpenSSH PEM format.
67+ privPEM , err := gossh .MarshalPrivateKey (priv , "" )
68+ if err != nil {
69+ return "" , "" , nil , fmt .Errorf ("could not encode private key: %w" , err )
70+ }
71+ privBytes := pem .EncodeToMemory (privPEM )
72+
73+ // Write private key to a temp file with restricted permissions.
74+ f , err := os .CreateTemp ("" , "createos-vm-*.pem" )
75+ if err != nil {
76+ return "" , "" , nil , fmt .Errorf ("could not create temp key file: %w" , err )
77+ }
78+ if err := os .Chmod (f .Name (), 0600 ); err != nil {
79+ _ = f .Close ()
80+ _ = os .Remove (f .Name ())
81+ return "" , "" , nil , fmt .Errorf ("could not set key file permissions: %w" , err )
82+ }
83+ if _ , err := f .Write (privBytes ); err != nil {
84+ _ = f .Close ()
85+ _ = os .Remove (f .Name ())
86+ return "" , "" , nil , fmt .Errorf ("could not write private key: %w" , err )
5287 }
88+ _ = f .Close ()
5389
54- const skipOption = "None (skip)"
55- options := make ([]string , 0 , len (keys )+ 1 )
90+ cleanup = func () { _ = os .Remove (f .Name ()) }
91+ return pubKeyStr , f .Name (), cleanup , nil
92+ }
93+
94+ const (
95+ optionTempKey = "Generate a temporary key (auto-deleted on exit)"
96+ optionSkip = "None (skip)"
97+ )
98+
99+ // selectPublicKey prompts the user to pick a local SSH public key, generate a temp
100+ // key, or skip. Returns the public key content and an optional private key path
101+ // (non-empty only for generated temp keys).
102+ func selectPublicKey () (publicKey string , privateKeyPath string , err error ) {
103+ keys := findPublicKeys ()
104+
105+ options := make ([]string , 0 , len (keys )+ 2 )
56106 nameByOption := make (map [string ]string , len (keys ))
57107
58108 for name , content := range keys {
59- // Show filename + key comment (last field) for readability.
60109 parts := strings .Fields (content )
61110 label := name
62111 if len (parts ) >= 3 {
@@ -65,19 +114,28 @@ func selectPublicKey() (string, error) {
65114 options = append (options , label )
66115 nameByOption [label ] = content
67116 }
68- options = append (options , skipOption )
117+ options = append (options , optionTempKey , optionSkip )
69118
70119 selected , err := pterm .DefaultInteractiveSelect .
71120 WithOptions (options ).
72121 WithDefaultText ("Select an SSH public key to use for this session" ).
73122 Show ()
74123 if err != nil {
75- return "" , fmt .Errorf ("could not read selection: %w" , err )
124+ return "" , "" , fmt .Errorf ("could not read selection: %w" , err )
76125 }
77- if selected == skipOption {
78- return "" , nil
126+
127+ switch selected {
128+ case optionSkip :
129+ return "" , "" , nil
130+ case optionTempKey :
131+ pub , privPath , _ , genErr := generateTempSSHKeypair ()
132+ if genErr != nil {
133+ return "" , "" , genErr
134+ }
135+ return pub , privPath , nil
136+ default :
137+ return nameByOption [selected ], "" , nil
79138 }
80- return nameByOption [selected ], nil
81139}
82140
83141func newVMSSHCommand () * cli.Command {
@@ -122,11 +180,16 @@ func newVMSSHCommand() *cli.Command {
122180 return fmt .Errorf ("VM does not have an IP address yet — try again in a moment\n \n Check status with: createos vms get %s" , id )
123181 }
124182
125- // Ask user to pick a local SSH public key for this session.
126- localKey , err := selectPublicKey ()
183+ localKey , privateKeyPath , err := selectPublicKey ()
127184 if err != nil {
128185 return err
129186 }
187+
188+ // Clean up temp private key file on exit if one was generated.
189+ if privateKeyPath != "" {
190+ defer os .Remove (privateKeyPath ) //nolint:errcheck
191+ }
192+
130193 originalKeys := vm .Inputs .SSHKeys
131194 if originalKeys == nil {
132195 originalKeys = []string {}
@@ -137,7 +200,6 @@ func newVMSSHCommand() *cli.Command {
137200 }
138201
139202 if localKey != "" {
140- // Check if key is already present.
141203 alreadyPresent := false
142204 for _ , k := range originalKeys {
143205 if strings .TrimSpace (k ) == localKey {
@@ -149,10 +211,9 @@ func newVMSSHCommand() *cli.Command {
149211 if ! alreadyPresent {
150212 updatedKeys := append (originalKeys , localKey ) //nolint:gocritic
151213 if err := client .UpdateVMDeployment (id , updatedKeys , firewallRules ); err != nil {
152- pterm .Warning .Println ("Could not add your SSH key to the VM — you may need to add it manually." )
214+ pterm .Warning .Printf ("Could not add your SSH key to the VM: %v \n " , err )
153215 } else {
154216 pterm .Info .Println ("Your SSH public key has been added to the VM." )
155- // Restore original keys on exit.
156217 defer func () {
157218 if err := client .UpdateVMDeployment (id , originalKeys , firewallRules ); err != nil {
158219 pterm .Warning .Printf ("Could not remove your SSH key from VM %q: %v\n " , id , err )
@@ -165,7 +226,13 @@ func newVMSSHCommand() *cli.Command {
165226 }
166227
167228 target := user + "@" + vm .Extra .IPAddress
168- cmd := exec .CommandContext (context .Background (), "ssh" , "-o" , "StrictHostKeyChecking=accept-new" , target ) //nolint:gosec // target is user@<api-provided-ip>, intentional subprocess
229+ sshArgs := []string {"-o" , "StrictHostKeyChecking=accept-new" }
230+ if privateKeyPath != "" {
231+ sshArgs = append (sshArgs , "-i" , privateKeyPath )
232+ }
233+ sshArgs = append (sshArgs , target )
234+
235+ cmd := exec .CommandContext (context .Background (), "ssh" , sshArgs ... ) //nolint:gosec
169236 cmd .Stdin = os .Stdin
170237 cmd .Stdout = os .Stdout
171238 cmd .Stderr = os .Stderr
0 commit comments