Repository: cli/cli
Pull request #5752 (2022-06-03 18:44:53)

Automatically create ssh keys in `gh cs ssh`
Fixes: https://github.com/github/codespaces/issues/8154

Adds automatic ssh key generation to `gh cs ssh` and `gh cs cp` commands, effectively removing the need for the user to already have set up ssh keys in GitHub.

The key is named `codespaces` and is stored under `~/.ssh`. If it already exists it is just reused, otherwise it is generated without prompt and without passphrase. When starting the ssh server on the codespace (over the RPC call), an additional argument is sent, which is the public key. The codespace agent will append this public key to the `authorized_keys` file.

This functionality requires a change in the codespaces agent to support, so it should NOT be merged until that change is released.

![ssh](https://user-images.githubusercontent.com/5447118/172884886-7b1fd273-3b70-4a30-8ed0-e3d9b7fa5118.gif)


Patch:
diff --git a/pkg/cmd/auth/shared/login_flow.go b/pkg/cmd/auth/shared/login_flow.go
index f2c0fd6fb..48b0c773c 100644
--- a/pkg/cmd/auth/shared/login_flow.go
+++ b/pkg/cmd/auth/shared/login_flow.go
@@ -3,6 +3,7 @@ package shared
 import (
 	"fmt"
 	"net/http"
+	"os"
 	"strings"
 
 	"github.com/AlecAivazis/survey/v2"
@@ -10,8 +11,10 @@ import (
 	"github.com/cli/cli/v2/api"
 	"github.com/cli/cli/v2/internal/authflow"
 	"github.com/cli/cli/v2/internal/ghinstance"
+	"github.com/cli/cli/v2/pkg/cmd/ssh-key/add"
 	"github.com/cli/cli/v2/pkg/iostreams"
 	"github.com/cli/cli/v2/pkg/prompt"
+	"github.com/cli/cli/v2/pkg/ssh"
 )
 
 const defaultSSHKeyTitle = "GitHub CLI"
@@ -34,7 +37,7 @@ type LoginOptions struct {
 	Executable  string
 	GitProtocol string
 
-	sshContext sshContext
+	sshContext ssh.Context
 }
 
 func Login(opts *LoginOptions) error {
@@ -72,7 +75,7 @@ func Login(opts *LoginOptions) error {
 	var keyToUpload string
 	keyTitle := defaultSSHKeyTitle
 	if opts.Interactive && gitProtocol == "ssh" {
-		pubKeys, err := opts.sshContext.localPublicKeys()
+		pubKeys, err := opts.sshContext.LocalPublicKeys()
 		if err != nil {
 			return err
 		}
@@ -89,11 +92,28 @@ func Login(opts *LoginOptions) error {
 			if keyChoice < len(pubKeys) {
 				keyToUpload = pubKeys[keyChoice]
 			}
-		} else {
-			var err error
-			keyToUpload, err = opts.sshContext.generateSSHKey()
+		} else if opts.sshContext.HasKeygen() {
+			var sshChoice bool
+			err := prompt.SurveyAskOne(&survey.Confirm{
+				Message: "Generate a new SSH key to add to your GitHub account?",
+				Default: true,
+			}, &sshChoice)
+
 			if err != nil {
-				return err
+				return fmt.Errorf("could not prompt: %w", err)
+			}
+
+			if sshChoice {
+				passphrase, err := promptForSshKeyPassphrase()
+				if err != nil {
+					return fmt.Errorf("could not prompt for key passphrase: %w", err)
+				}
+
+				keyPair, err := opts.sshContext.GenerateSSHKey("id_ed25519", passphrase)
+				if err != nil {
+					return err
+				}
+				keyToUpload = keyPair.PublicKeyPath
 			}
 		}
 
@@ -210,6 +230,19 @@ func Login(opts *LoginOptions) error {
 	return nil
 }
 
+func promptForSshKeyPassphrase() (string, error) {
+	var sshPassphrase string
+	err := prompt.SurveyAskOne(&survey.Password{
+		Message: "Enter a passphrase for your new SSH key (Optional)",
+	}, &sshPassphrase)
+
+	if err != nil {
+		return "", err
+	}
+
+	return sshPassphrase, nil
+}
+
 func scopesSentence(scopes []string, isEnterprise bool) string {
 	quoted := make([]string, len(scopes))
 	for i, s := range scopes {
@@ -221,3 +254,13 @@ func scopesSentence(scopes []string, isEnterprise bool) string {
 	}
 	return strings.Join(quoted, ", ")
 }
+
+func sshKeyUpload(httpClient *http.Client, hostname, keyFile string, title string) error {
+	f, err := os.Open(keyFile)
+	if err != nil {
+		return err
+	}
+	defer f.Close()
+
+	return add.SSHKeyUpload(httpClient, hostname, f, title)
+}
diff --git a/pkg/cmd/auth/shared/ssh_keys.go b/pkg/cmd/auth/shared/ssh_keys.go
deleted file mode 100644
index 4a8914a2c..000000000
--- a/pkg/cmd/auth/shared/ssh_keys.go
+++ /dev/null
@@ -1,119 +0,0 @@
-package shared
-
-import (
-	"fmt"
-	"net/http"
-	"os"
-	"os/exec"
-	"path/filepath"
-	"runtime"
-
-	"github.com/AlecAivazis/survey/v2"
-	"github.com/cli/cli/v2/internal/config"
-	"github.com/cli/cli/v2/internal/run"
-	"github.com/cli/cli/v2/pkg/cmd/ssh-key/add"
-	"github.com/cli/cli/v2/pkg/prompt"
-	"github.com/cli/safeexec"
-)
-
-type sshContext struct {
-	configDir string
-	keygenExe string
-}
-
-func (c *sshContext) sshDir() (string, error) {
-	if c.configDir != "" {
-		return c.configDir, nil
-	}
-	dir, err := config.HomeDirPath(".ssh")
-	if err == nil {
-		c.configDir = dir
-	}
-	return dir, err
-}
-
-func (c *sshContext) localPublicKeys() ([]string, error) {
-	sshDir, err := c.sshDir()
-	if err != nil {
-		return nil, err
-	}
-
-	return filepath.Glob(filepath.Join(sshDir, "*.pub"))
-}
-
-func (c *sshContext) findKeygen() (string, error) {
-	if c.keygenExe != "" {
-		return c.keygenExe, nil
-	}
-
-	keygenExe, err := safeexec.LookPath("ssh-keygen")
-	if err != nil && runtime.GOOS == "windows" {
-		// We can try and find ssh-keygen in a Git for Windows install
-		if gitPath, err := safeexec.LookPath("git"); err == nil {
-			gitKeygen := filepath.Join(filepath.Dir(gitPath), "..", "usr", "bin", "ssh-keygen.exe")
-			if _, err = os.Stat(gitKeygen); err == nil {
-				return gitKeygen, nil
-			}
-		}
-	}
-
-	if err == nil {
-		c.keygenExe = keygenExe
-	}
-	return keygenExe, err
-}
-
-func (c *sshContext) generateSSHKey() (string, error) {
-	keygenExe, err := c.findKeygen()
-	if err != nil {
-		// give up silently if `ssh-keygen` is not available
-		return "", nil
-	}
-
-	var sshChoice bool
-	err = prompt.SurveyAskOne(&survey.Confirm{
-		Message: "Generate a new SSH key to add to your GitHub account?",
-		Default: true,
-	}, &sshChoice)
-	if err != nil {
-		return "", fmt.Errorf("could not prompt: %w", err)
-	}
-	if !sshChoice {
-		return "", nil
-	}
-
-	sshDir, err := c.sshDir()
-	if err != nil {
-		return "", err
-	}
-	keyFile := filepath.Join(sshDir, "id_ed25519")
-	if _, err := os.Stat(keyFile); err == nil {
-		return "", fmt.Errorf("refusing to overwrite file %s", keyFile)
-	}
-
-	if err := os.MkdirAll(filepath.Dir(keyFile), 0711); err != nil {
-		return "", err
-	}
-
-	var sshLabel string
-	var sshPassphrase string
-	err = prompt.SurveyAskOne(&survey.Password{
-		Message: "Enter a passphrase for your new SSH key (Optional)",
-	}, &sshPassphrase)
-	if err != nil {
-		return "", fmt.Errorf("could not prompt: %w", err)
-	}
-
-	keygenCmd := exec.Command(keygenExe, "-t", "ed25519", "-C", sshLabel, "-N", sshPassphrase, "-f", keyFile)
-	return keyFile + ".pub", run.PrepareCmd(keygenCmd).Run()
-}
-
-func sshKeyUpload(httpClient *http.Client, hostname, keyFile string, title string) error {
-	f, err := os.Open(keyFile)
-	if err != nil {
-		return err
-	}
-	defer f.Close()
-
-	return add.SSHKeyUpload(httpClient, hostname, f, title)
-}
diff --git a/pkg/cmd/codespace/common.go b/pkg/cmd/codespace/common.go
index 402fcd7ea..6720aa0bb 100644
--- a/pkg/cmd/codespace/common.go
+++ b/pkg/cmd/codespace/common.go
@@ -69,6 +69,7 @@ type liveshareSession interface {
 	StartJupyterServer(context.Context) (int, string, error)
 	StartSharing(context.Context, string, int) (liveshare.ChannelID, error)
 	StartSSHServer(context.Context) (int, string, error)
+	StartSSHServerWithOptions(context.Context, liveshare.StartSSHServerOptions) (int, string, error)
 }
 
 // Connects to a codespace using Live Share and returns that session
diff --git a/pkg/cmd/codespace/ssh.go b/pkg/cmd/codespace/ssh.go
index 10ea30473..2c3f1c93f 100644
--- a/pkg/cmd/codespace/ssh.go
+++ b/pkg/cmd/codespace/ssh.go
@@ -20,6 +20,7 @@ import (
 	"github.com/cli/cli/v2/internal/codespaces/api"
 	"github.com/cli/cli/v2/pkg/cmdutil"
 	"github.com/cli/cli/v2/pkg/liveshare"
+	"github.com/cli/cli/v2/pkg/ssh"
 	"github.com/spf13/cobra"
 )
 
@@ -115,6 +116,25 @@ func (a *App) SSH(ctx context.Context, sshArgs []string, opts sshOptions) (err e
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
+	args := sshArgs
+	if opts.scpArgs != nil {
+		args = opts.scpArgs
+	}
+
+	sshContext := ssh.Context{}
+	startSSHOptions := liveshare.StartSSHServerOptions{}
+
+	if shouldGenerateSSHKeys(args, opts) && sshContext.HasKeygen() {
+		keyPair, err := sshContext.GenerateSSHKey("codespaces", "")
+		if err != nil && !errors.Is(err, ssh.ErrKeyAlreadyExists) {
+			return fmt.Errorf("failed to generate ssh keys: %w", err)
+		}
+
+		startSSHOptions.UserPublicKeyFile = keyPair.PublicKeyPath
+		// For both cp and ssh, flags need to come first in the args (before a command in ssh and files in cp), so prepend this flag
+		args = append([]string{"-i", keyPair.PrivateKeyPath}, args...)
+	}
+
 	codespace, err := getOrChooseCodespace(ctx, a.apiClient, opts.codespace)
 	if err != nil {
 		return err
@@ -127,7 +147,7 @@ func (a *App) SSH(ctx context.Context, sshArgs []string, opts sshOptions) (err e
 	defer safeClose(session, &err)
 
 	a.StartProgressIndicatorWithLabel("Fetching SSH Details")
-	remoteSSHServerPort, sshUser, err := session.StartSSHServer(ctx)
+	remoteSSHServerPort, sshUser, err := session.StartSSHServerWithOptions(ctx, startSSHOptions)
 	a.StopProgressIndicator()
 	if err != nil {
 		return fmt.Errorf("error getting ssh server details: %w", err)
@@ -168,9 +188,10 @@ func (a *App) SSH(ctx context.Context, sshArgs []string, opts sshOptions) (err e
 	go func() {
 		var err error
 		if opts.scpArgs != nil {
-			err = codespaces.Copy(ctx, opts.scpArgs, localSSHServerPort, connectDestination)
+			// args is the correct variable to use here, we just use scpArgs as the check for which command to run
+			err = codespaces.Copy(ctx, args, localSSHServerPort, connectDestination)
 		} else {
-			err = codespaces.Shell(ctx, a.errLogger, sshArgs, localSSHServerPort, connectDestination, usingCustomPort)
+			err = codespaces.Shell(ctx, a.errLogger, args, localSSHServerPort, connectDestination, usingCustomPort)
 		}
 		shellClosed <- err
 	}()
@@ -186,6 +207,27 @@ func (a *App) SSH(ctx context.Context, sshArgs []string, opts sshOptions) (err e
 	}
 }
 
+func shouldGenerateSSHKeys(args []string, opts sshOptions) bool {
+	if opts.profile != "" {
+		// The profile may specify the identity file so cautiously don't override anything with that option
+		return false
+	}
+
+	for _, arg := range args {
+		if arg == "-i" {
+			// User specified the identity file so it should exist
+			return false
+		}
+
+		if arg == "-F" {
+			// User specified a config file so, similar to profile, cautiously don't override settings
+			return false
+		}
+	}
+
+	return true
+}
+
 func (a *App) printOpenSSHConfig(ctx context.Context, opts sshOptions) (err error) {
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
diff --git a/pkg/liveshare/session.go b/pkg/liveshare/session.go
index 33531bffd..f912fa5ea 100644
--- a/pkg/liveshare/session.go
+++ b/pkg/liveshare/session.go
@@ -3,7 +3,9 @@ package liveshare
 import (
 	"context"
 	"fmt"
+	"os"
 	"strconv"
+	"strings"
 	"time"
 
 	"github.com/opentracing/opentracing-go"
@@ -27,6 +29,10 @@ type Session struct {
 	logger          logger
 }
 
+type StartSSHServerOptions struct {
+	UserPublicKeyFile string
+}
+
 // Close should be called by users to clean up RPC and SSH resources whenever the session
 // is no longer active.
 func (s *Session) Close() error {
@@ -46,9 +52,19 @@ func (s *Session) registerRequestHandler(requestType string, h handler) func() {
 	return s.rpc.register(requestType, h)
 }
 
-// StartsSSHServer starts an SSH server in the container, installing sshd if necessary,
-// and returns the port on which it listens and the user name clients should provide.
+// StartSSHServer starts an SSH server in the container, installing sshd if necessary, applies specified
+// options, and returns the port on which it listens and the user name clients should provide.
 func (s *Session) StartSSHServer(ctx context.Context) (int, string, error) {
+	return s.StartSSHServerWithOptions(ctx, StartSSHServerOptions{})
+}
+
+// StartSSHServerWithOptions starts an SSH server in the container, installing sshd if necessary, applies specified
+// options, and returns the port on which it listens and the user name clients should provide.
+func (s *Session) StartSSHServerWithOptions(ctx context.Context, options StartSSHServerOptions) (int, string, error) {
+	var params struct {
+		UserPublicKey string `json:"userPublicKey"`
+	}
+
 	var response struct {
 		Result     bool   `json:"result"`
 		ServerPort string `json:"serverPort"`
@@ -56,7 +72,16 @@ func (s *Session) StartSSHServer(ctx context.Context) (int, string, error) {
 		Message    string `json:"message"`
 	}
 
-	if err := s.rpc.do(ctx, "ISshServerHostService.startRemoteServer", []string{}, &response); err != nil {
+	if options.UserPublicKeyFile != "" {
+		publicKeyBytes, err := os.ReadFile(options.UserPublicKeyFile)
+		if err != nil {
+			return 0, "", fmt.Errorf("failed to read public key file: %w", err)
+		}
+
+		params.UserPublicKey = strings.TrimSpace(string(publicKeyBytes))
+	}
+
+	if err := s.rpc.do(ctx, "ISshServerHostService.startRemoteServerWithOptions", params, &response); err != nil {
 		return 0, "", err
 	}
 
diff --git a/pkg/ssh/ssh_keys.go b/pkg/ssh/ssh_keys.go
new file mode 100644
index 000000000..c750b608a
--- /dev/null
+++ b/pkg/ssh/ssh_keys.go
@@ -0,0 +1,113 @@
+package ssh
+
+import (
+	"errors"
+	"fmt"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"runtime"
+
+	"github.com/cli/cli/v2/internal/config"
+	"github.com/cli/cli/v2/internal/run"
+	"github.com/cli/safeexec"
+)
+
+type Context struct {
+	ConfigDir string
+	KeygenExe string
+}
+
+type KeyPair struct {
+	PublicKeyPath  string
+	PrivateKeyPath string
+}
+
+var ErrKeyAlreadyExists = errors.New("SSH key already exists")
+
+func (c *Context) LocalPublicKeys() ([]string, error) {
+	sshDir, err := c.sshDir()
+	if err != nil {
+		return nil, err
+	}
+
+	return filepath.Glob(filepath.Join(sshDir, "*.pub"))
+}
+
+func (c *Context) HasKeygen() bool {
+	_, err := c.findKeygen()
+	return err == nil
+}
+
+func (c *Context) GenerateSSHKey(keyName string, passphrase string) (*KeyPair, error) {
+	keygenExe, err := c.findKeygen()
+	if err != nil {
+		return nil, err
+	}
+
+	sshDir, err := c.sshDir()
+	if err != nil {
+		return nil, err
+	}
+
+	err = os.MkdirAll(sshDir, 0700)
+	if err != nil {
+		return nil, fmt.Errorf("could not create .ssh directory: %w", err)
+	}
+
+	keyFile := filepath.Join(sshDir, keyName)
+	keyPair := KeyPair{
+		PublicKeyPath:  keyFile + ".pub",
+		PrivateKeyPath: keyFile,
+	}
+
+	if _, err := os.Stat(keyFile); err == nil {
+		// Still return keyPair because the caller might be OK with this - they can check the error with errors.Is(err, ErrKeyAlreadyExists)
+		return &keyPair, ErrKeyAlreadyExists
+	}
+
+	if err := os.MkdirAll(filepath.Dir(keyFile), 0711); err != nil {
+		return nil, err
+	}
+
+	keygenCmd := exec.Command(keygenExe, "-t", "ed25519", "-C", "", "-N", passphrase, "-f", keyFile)
+	err = run.PrepareCmd(keygenCmd).Run()
+	if err != nil {
+		return nil, err
+	}
+
+	return &keyPair, nil
+}
+
+func (c *Context) sshDir() (string, error) {
+	if c.ConfigDir != "" {
+		return c.ConfigDir, nil
+	}
+	dir, err := config.HomeDirPath(".ssh")
+	if err == nil {
+		c.ConfigDir = dir
+	}
+	return dir, err
+}
+
+func (c *Context) findKeygen() (string, error) {
+	if c.KeygenExe != "" {
+		return c.KeygenExe, nil
+	}
+
+	keygenExe, err := safeexec.LookPath("ssh-keygen")
+	if err != nil && runtime.GOOS == "windows" {
+		// We can try and find ssh-keygen in a Git for Windows install
+		if gitPath, err := safeexec.LookPath("git"); err == nil {
+			gitKeygen := filepath.Join(filepath.Dir(gitPath), "..", "usr", "bin", "ssh-keygen.exe")
+			if _, err = os.Stat(gitKeygen); err == nil {
+				return gitKeygen, nil
+			}
+		}
+	}
+
+	if err == nil {
+		c.KeygenExe = keygenExe
+	}
+	return keygenExe, err
+}


Test patch:
diff --git a/pkg/cmd/auth/shared/login_flow_test.go b/pkg/cmd/auth/shared/login_flow_test.go
index 60a19fe03..d32e2cc87 100644
--- a/pkg/cmd/auth/shared/login_flow_test.go
+++ b/pkg/cmd/auth/shared/login_flow_test.go
@@ -12,6 +12,7 @@ import (
 	"github.com/cli/cli/v2/pkg/httpmock"
 	"github.com/cli/cli/v2/pkg/iostreams"
 	"github.com/cli/cli/v2/pkg/prompt"
+	"github.com/cli/cli/v2/pkg/ssh"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -84,9 +85,9 @@ func TestLogin_ssh(t *testing.T) {
 		HTTPClient:  &http.Client{Transport: &tr},
 		Hostname:    "example.com",
 		Interactive: true,
-		sshContext: sshContext{
-			configDir: dir,
-			keygenExe: "ssh-keygen",
+		sshContext: ssh.Context{
+			ConfigDir: dir,
+			KeygenExe: "ssh-keygen",
 		},
 	})
 	assert.NoError(t, err)