Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
72 changes: 61 additions & 11 deletions cmd/ssh.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ package cmd
import (
"context"
"encoding/base64"
"encoding/json"
"fmt"
"os"
"os/exec"
Expand All @@ -16,6 +17,15 @@ import (
"github.com/spf13/cobra"
)

// sshSetupResult is the JSON output structure for --setup-only -o json.
type sshSetupResult struct {
VMDomain string `json:"vm_domain"`
SessionID string `json:"session_id"`
SSHKeyFile string `json:"ssh_key_file"`
ProxyCommand string `json:"proxy_command"`
SSHCommand string `json:"ssh_command"`
}

var sshCmd = &cobra.Command{
Use: "ssh <id>",
Short: "Open an interactive SSH session to a browser VM",
Expand Down Expand Up @@ -49,6 +59,7 @@ func init() {
sshCmd.Flags().StringP("local-forward", "L", "", "Local port forwarding (localport:host:remoteport)")
sshCmd.Flags().StringP("remote-forward", "R", "", "Remote port forwarding (remoteport:host:localport)")
sshCmd.Flags().Bool("setup-only", false, "Setup SSH on VM without connecting")
sshCmd.Flags().StringP("output", "o", "", "Output format: json for machine-readable output (only with --setup-only)")
}

func runSSH(cmd *cobra.Command, args []string) error {
Expand All @@ -60,26 +71,39 @@ func runSSH(cmd *cobra.Command, args []string) error {
localForward, _ := cmd.Flags().GetString("local-forward")
remoteForward, _ := cmd.Flags().GetString("remote-forward")
setupOnly, _ := cmd.Flags().GetBool("setup-only")
output, _ := cmd.Flags().GetString("output")

if output != "" && output != "json" {
return fmt.Errorf("unsupported --output value: use 'json'")
}
if output == "json" && !setupOnly {
return fmt.Errorf("--output json is only supported with --setup-only")
}

cfg := ssh.Config{
BrowserID: browserID,
IdentityFile: identityFile,
LocalForward: localForward,
RemoteForward: remoteForward,
SetupOnly: setupOnly,
Output: output,
}

return connectSSH(ctx, client, cfg)
}

func connectSSH(ctx context.Context, client kernel.Client, cfg ssh.Config) error {
jsonOutput := cfg.Output == "json"

// Check websocat is installed locally
if err := ssh.CheckWebsocatInstalled(); err != nil {
return err
}

// Get browser info
pterm.Info.Printf("Getting browser %s info...\n", cfg.BrowserID)
if !jsonOutput {
pterm.Info.Printf("Getting browser %s info...\n", cfg.BrowserID)
}
browser, err := client.Browsers.Get(ctx, cfg.BrowserID, kernel.BrowserGetParams{})
if err != nil {
return fmt.Errorf("failed to get browser: %w", err)
Expand All @@ -95,7 +119,9 @@ func connectSSH(ctx context.Context, client kernel.Client, cfg ssh.Config) error
if err != nil {
return fmt.Errorf("failed to extract VM domain: %w", err)
}
pterm.Info.Printf("VM domain: %s\n", vmDomain)
if !jsonOutput {
pterm.Info.Printf("VM domain: %s\n", vmDomain)
}

// Generate or load SSH keypair
var privateKeyPEM, publicKey string
Expand All @@ -104,7 +130,9 @@ func connectSSH(ctx context.Context, client kernel.Client, cfg ssh.Config) error

if cfg.IdentityFile != "" {
// Use provided key
pterm.Info.Printf("Using SSH key: %s\n", cfg.IdentityFile)
if !jsonOutput {
pterm.Info.Printf("Using SSH key: %s\n", cfg.IdentityFile)
}
keyFile = cfg.IdentityFile

// Read public key to inject into VM
Expand All @@ -117,7 +145,9 @@ func connectSSH(ctx context.Context, client kernel.Client, cfg ssh.Config) error
publicKey = strings.TrimSpace(string(pubKeyData))
} else {
// Generate ephemeral keypair
pterm.Info.Println("Generating ephemeral SSH keypair...")
if !jsonOutput {
pterm.Info.Println("Generating ephemeral SSH keypair...")
}
keyPair, err := ssh.GenerateKeyPair()
if err != nil {
return fmt.Errorf("failed to generate SSH keypair: %w", err)
Expand All @@ -134,22 +164,40 @@ func connectSSH(ctx context.Context, client kernel.Client, cfg ssh.Config) error
pterm.Debug.Printf("Temp key file: %s\n", keyFile)
}

// Cleanup temp key on exit
if cleanupKey {
// Cleanup temp key on exit (skip if JSON setup-only, since the caller needs the key)
if cleanupKey && !(cfg.SetupOnly && jsonOutput) {
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ephemeral key cleanup condition too narrow for setup-only

Medium Severity

The cleanup guard !(cfg.SetupOnly && jsonOutput) only preserves the ephemeral key file when JSON output is requested. In the non-JSON --setup-only path, the key is still deleted by the deferred os.Remove before the user can use the printed manual SSH command. The condition recognizes "the caller needs the key" (per the comment) but the guard needs to be !cfg.SetupOnly to cover both output modes.

Additional Locations (1)

Fix in Cursor Fix in Web

defer func() {
pterm.Debug.Printf("Cleaning up temp key: %s\n", keyFile)
os.Remove(keyFile)
}()
}

// Setup SSH services on VM
pterm.Info.Println("Setting up SSH services on VM...")
if err := setupVMSSH(ctx, client, browser.SessionID, publicKey); err != nil {
if !jsonOutput {
pterm.Info.Println("Setting up SSH services on VM...")
}
if err := setupVMSSH(ctx, client, browser.SessionID, publicKey, jsonOutput); err != nil {
return fmt.Errorf("failed to setup SSH on VM: %w", err)
}
pterm.Success.Println("SSH services running on VM")
if !jsonOutput {
pterm.Success.Println("SSH services running on VM")
}

if cfg.SetupOnly {
if jsonOutput {
proxyCmd := fmt.Sprintf("websocat --binary wss://%s:2222", vmDomain)
sshCommand := fmt.Sprintf("ssh -o 'ProxyCommand=%s' -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -o LogLevel=ERROR -i %s root@localhost", proxyCmd, keyFile)
result := sshSetupResult{
VMDomain: vmDomain,
SessionID: browser.SessionID,
SSHKeyFile: keyFile,
ProxyCommand: proxyCmd,
SSHCommand: sshCommand,
}
enc := json.NewEncoder(os.Stdout)
enc.SetIndent("", " ")
return enc.Encode(result)
}
pterm.Info.Println("\n--setup-only specified, not connecting.")
pterm.Info.Printf("To connect manually:\n")
pterm.Info.Printf(" ssh -o 'ProxyCommand=websocat --binary wss://%s:2222' -i %s root@localhost\n", vmDomain, keyFile)
Expand Down Expand Up @@ -192,7 +240,7 @@ func connectSSH(ctx context.Context, client kernel.Client, cfg ssh.Config) error
}

// setupVMSSH installs and configures sshd + websocat on the VM using process.exec
func setupVMSSH(ctx context.Context, client kernel.Client, sessionID, publicKey string) error {
func setupVMSSH(ctx context.Context, client kernel.Client, sessionID, publicKey string, quiet bool) error {
// First check if services are already running
checkScript := ssh.CheckServicesScript()
checkResp, err := client.Browsers.Process.Exec(ctx, sessionID, kernel.BrowserProcessExecParams{
Expand All @@ -205,7 +253,9 @@ func setupVMSSH(ctx context.Context, client kernel.Client, sessionID, publicKey
} else if checkResp != nil && checkResp.StdoutB64 != "" {
stdout, _ := base64.StdEncoding.DecodeString(checkResp.StdoutB64)
if strings.TrimSpace(string(stdout)) == "RUNNING" {
pterm.Info.Println("SSH services already running, injecting key...")
if !quiet {
pterm.Info.Println("SSH services already running, injecting key...")
}
// Just inject the key
return injectSSHKey(ctx, client, sessionID, publicKey)
}
Expand Down
1 change: 1 addition & 0 deletions pkg/ssh/ssh.go
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ type Config struct {
LocalForward string // -L flag value
RemoteForward string // -R flag value
SetupOnly bool
Output string // "json" for machine-readable output
}

// KeyPair holds an SSH keypair
Expand Down