add cmdexec package to execute commands on remote hosts via SSH

internal/cmdexec/executor.go

@@ -0,0 +1,41 @@
+package cmdexec
+
+import (
+	"context"
+	"regexp"
+	"strings"
+)
+
+type Executor interface {
+	Run(ctx context.Context, cmd string) (string, error)
+	Close() error
+}
+
+// Quote* functions are copied from github.com/alessio/shellescape package.
+var pattern = regexp.MustCompile(`[^\w@%+=:,./-]`)
+
+// Quote returns a shell-escaped version of the string s. The returned value
+// is a string that can safely be used as one token in a shell command line.
+func Quote(s string) string {
+	if len(s) == 0 {
+		return "''"
+	}
+
+	if pattern.MatchString(s) {
+		return "'" + strings.ReplaceAll(s, "'", "'\"'\"'") + "'"
+	}
+
+	return s
+}
+
+// QuoteCommand returns a shell-escaped version of the command arguments.
+// The returned value is a string that can safely be used as shell command arguments.
+func QuoteCommand(args ...string) string {
+	l := make([]string, len(args))
+
+	for i, s := range args {
+		l[i] = Quote(s)
+	}
+
+	return strings.Join(l, " ")
+}

internal/cmdexec/remote.go

@@ -0,0 +1,59 @@
+package cmdexec
+
+import (
+	"context"
+	"fmt"
+	"golang.org/x/crypto/ssh"
+	"strings"
+)
+
+type Remote struct {
+	client *ssh.Client
+}
+
+// Run runs the command on the remote host and returns its output with all leading and trailing
+// white space removed.
+func (r *Remote) Run(ctx context.Context, cmd string) (string, error) {
+	session, err := r.client.NewSession()
+	if err != nil {
+		return "", fmt.Errorf("create session: %w", err)
+	}
+	defer func() {
+		_ = session.Close()
+	}()
+
+	// Run the command in a goroutine to be able to cancel it.
+	type result struct {
+		out string
+		err error
+	}
+	done := make(chan result)
+	go func() {
+		outBytes, outErr := session.CombinedOutput(cmd)
+		res := result{
+			err: outErr,
+		}
+		if outErr == nil {
+			res.out = strings.TrimSpace(string(outBytes))
+		}
+		done <- res
+	}()
+
+	select {
+	case res := <-done:
+		if res.err != nil {
+			return "", fmt.Errorf("run command on remote host: %w", res.err)
+		}
+		return res.out, nil
+	case <-ctx.Done():
+		if err = session.Signal(ssh.SIGINT); err != nil {
+			return "", fmt.Errorf("send interrupt signal to remote process: %w", err)
+		}
+		return "", fmt.Errorf("canceled: %w", ctx.Err())
+	}
+}
+
+// Close closes the connection to the remote host.
+func (r *Remote) Close() error {
+	return r.client.Close()
+}

internal/cmdexec/ssh.go

@@ -0,0 +1,77 @@
+package cmdexec
+
+import (
+	"fmt"
+	"golang.org/x/crypto/ssh"
+	"golang.org/x/crypto/ssh/agent"
+	"net"
+	"os"
+	"strconv"
+	"time"
+)
+
+func Connect(user, host string, port int, sshKeyPath string) (*Remote, error) {
+	addr := net.JoinHostPort(host, strconv.Itoa(port))
+	// Try to connect using SSH agent only.
+	agentAuth, agentClose, agentErr := sshAgentAuth()
+	if agentErr == nil {
+		defer agentClose()
+		config := &ssh.ClientConfig{
+			User:            user,
+			Auth:            []ssh.AuthMethod{agentAuth},
+			HostKeyCallback: ssh.InsecureIgnoreHostKey(),
+			Timeout:         5 * time.Second,
+		}
+		var client *ssh.Client
+		if client, agentErr = ssh.Dial("tcp", addr, config); agentErr == nil {
+			return &Remote{
+				client: client,
+			}, nil
+		}
+	}
+	// Fall back to using private key as the connection attempt using SSH agent failed.
+	if sshKeyPath == "" {
+		// TODO: iterate over ~/.ssh/id_* and try to connect using each key.
+		return nil, fmt.Errorf("connect using SSH agent: %w", agentErr)
+	}
+
+	keyAuth, err := privateKeyAuth(sshKeyPath)
+	config := &ssh.ClientConfig{
+		User:            user,
+		Auth:            []ssh.AuthMethod{keyAuth},
+		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
+		Timeout:         5 * time.Second,
+	}
+	client, err := ssh.Dial("tcp", addr, config)
+	if err != nil {
+		return nil, fmt.Errorf("connect using private key %q: %w", sshKeyPath, err)
+	}
+
+	return &Remote{
+		client: client,
+	}, nil
+}
+
+func sshAgentAuth() (ssh.AuthMethod, func(), error) {
+	conn, err := net.Dial("unix", os.Getenv("SSH_AUTH_SOCK"))
+	if err != nil {
+		return nil, func() {}, fmt.Errorf("connect to SSH agent: %w", err)
+	}
+	auth := ssh.PublicKeysCallback(agent.NewClient(conn).Signers)
+	connClose := func() { _ = conn.Close() }
+	return auth, connClose, nil
+}
+
+func privateKeyAuth(path string) (ssh.AuthMethod, error) {
+	key, err := os.ReadFile(path)
+	if err != nil {
+		return nil, fmt.Errorf("read private key file %q: %w", path, err)
+	}
+	signer, err := ssh.ParsePrivateKey(key)
+	if err != nil {
+		return nil, fmt.Errorf("parse private key: %w", err)
+	}
+	// TODO: prompt password for the private key if needed.
+	//  Check: https://github.com/alexellis/k3sup/blob/master/cmd/install.go
+	return ssh.PublicKeys(signer), nil
+}