llmgw/internal/onboard/gossh.go

package onboard

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"net"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"

	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/agent"
	"golang.org/x/crypto/ssh/knownhosts"
)

// GoSSH is Remote through Go's own ssh client rather than the system's: what
// the setup screens use, because they ask for a password when there is no key,
// which the system ssh in batch mode cannot take.
//
// Authentication is whichever is given, tried in order: the running ssh
// agent, a key file, a password. The password is held in memory for the
// session and never written anywhere.
type GoSSH struct {
	Addr string
	Port int
	User string

	UseAgent bool
	KeyFile  string
	// KeyPassphrase opens an encrypted KeyFile.
	KeyPassphrase string
	Password      string

	// SudoPassword is asked for when the user cannot sudo without one. With
	// it, PrepareSudo lets the user sudo without a password until
	// RevokeSudo, so every script can use plain sudo as on any other node.
	SudoPassword string

	// KnownHosts is where host keys are kept; "" is ~/.ssh/known_hosts.
	KnownHosts string

	mu        sync.Mutex
	client    *ssh.Client
	sudoGrant bool
}

// Run runs script with sh on the machine.
func (g *GoSSH) Run(ctx context.Context, script string, stdin []byte) ([]byte, error) {
	c, err := g.dial(ctx)
	if err != nil {
		return nil, err
	}
	sess, err := c.NewSession()
	if err != nil {
		// A connection that died since: dial again once.
		g.drop()
		if c, err = g.dial(ctx); err != nil {
			return nil, err
		}
		if sess, err = c.NewSession(); err != nil {
			return nil, err
		}
	}
	defer sess.Close()
	if stdin != nil {
		sess.Stdin = bytes.NewReader(stdin)
	}
	var stdout, stderr bytes.Buffer
	sess.Stdout, sess.Stderr = &stdout, &stderr
	done := make(chan error, 1)
	go func() { done <- sess.Run("sh -c " + shellQuote(script)) }()
	select {
	case <-ctx.Done():
		_ = sess.Signal(ssh.SIGKILL)
		return stdout.Bytes(), ctx.Err()
	case err := <-done:
		if err != nil {
			return append(stdout.Bytes(), stderr.Bytes()...), err
		}
		return stdout.Bytes(), nil
	}
}

// Close ends the connection, taking back a sudo grant first.
func (g *GoSSH) Close() error {
	_ = g.RevokeSudo(context.Background())
	g.drop()
	return nil
}

func (g *GoSSH) drop() {
	g.mu.Lock()
	defer g.mu.Unlock()
	if g.client != nil {
		_ = g.client.Close()
		g.client = nil
	}
}

func (g *GoSSH) target() string {
	port := g.Port
	if port == 0 {
		port = 22
	}
	return net.JoinHostPort(g.Addr, strconv.Itoa(port))
}

func (g *GoSSH) dial(ctx context.Context) (*ssh.Client, error) {
	g.mu.Lock()
	defer g.mu.Unlock()
	if g.client != nil {
		return g.client, nil
	}
	var methods []ssh.AuthMethod
	if g.UseAgent {
		if sock := os.Getenv("SSH_AUTH_SOCK"); sock != "" {
			if conn, err := net.Dial("unix", sock); err == nil {
				methods = append(methods, ssh.PublicKeysCallback(agent.NewClient(conn).Signers))
			}
		}
	}
	if g.KeyFile != "" {
		signer, err := loadKey(g.KeyFile, g.KeyPassphrase)
		if err != nil {
			return nil, err
		}
		methods = append(methods, ssh.PublicKeys(signer))
	}
	if g.Password != "" {
		pw := g.Password
		methods = append(methods, ssh.Password(pw),
			ssh.KeyboardInteractive(func(_, _ string, qs []string, _ []bool) ([]string, error) {
				out := make([]string, len(qs))
				for i := range qs {
					out[i] = pw
				}
				return out, nil
			}))
	}
	if len(methods) == 0 {
		return nil, fmt.Errorf("no way to log in to %s: no ssh agent, key or password", g.Addr)
	}
	hostKey, err := g.hostKeyCallback()
	if err != nil {
		return nil, err
	}
	cfg := &ssh.ClientConfig{User: g.User, Auth: methods, HostKeyCallback: hostKey, Timeout: 10 * time.Second}
	d := net.Dialer{Timeout: 10 * time.Second}
	conn, err := d.DialContext(ctx, "tcp", g.target())
	if err != nil {
		return nil, fmt.Errorf("cannot reach %s: %w", g.target(), err)
	}
	cc, chans, reqs, err := ssh.NewClientConn(conn, g.target(), cfg)
	if err != nil {
		_ = conn.Close()
		return nil, fmt.Errorf("cannot log in to %s as %s: %w", g.Addr, g.User, err)
	}
	g.client = ssh.NewClient(cc, chans, reqs)
	return g.client, nil
}

func loadKey(path, passphrase string) (ssh.Signer, error) {
	b, err := os.ReadFile(path)
	if err != nil {
		return nil, fmt.Errorf("reading the key %s: %w", path, err)
	}
	if passphrase != "" {
		return ssh.ParsePrivateKeyWithPassphrase(b, []byte(passphrase))
	}
	s, err := ssh.ParsePrivateKey(b)
	var missing *ssh.PassphraseMissingError
	if errors.As(err, &missing) {
		return nil, fmt.Errorf("the key %s needs its passphrase", path)
	}
	return s, err
}

// hostKeyCallback trusts a machine on first sight and records its key, and
// refuses one whose key has changed since — accept-new, as onboarding's
// system ssh does: a machine provisioned a minute ago has a key nobody has
// seen, and the risk guarded against is a machine substituted for one that
// was known.
func (g *GoSSH) hostKeyCallback() (ssh.HostKeyCallback, error) {
	path := g.KnownHosts
	if path == "" {
		home, err := os.UserHomeDir()
		if err != nil {
			return nil, err
		}
		path = filepath.Join(home, ".ssh", "known_hosts")
	}
	if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
		return nil, err
	}
	if f, err := os.OpenFile(path, os.O_CREATE|os.O_RDONLY, 0o600); err == nil {
		_ = f.Close()
	}
	known, err := knownhosts.New(path)
	if err != nil {
		return nil, fmt.Errorf("reading %s: %w", path, err)
	}
	return func(host string, remote net.Addr, key ssh.PublicKey) error {
		err := known(host, remote, key)
		var ke *knownhosts.KeyError
		if errors.As(err, &ke) && len(ke.Want) == 0 {
			f, ferr := os.OpenFile(path, os.O_APPEND|os.O_WRONLY, 0o600)
			if ferr != nil {
				return ferr
			}
			defer f.Close()
			_, ferr = fmt.Fprintln(f, knownhosts.Line([]string{knownhosts.Normalize(host)}, key))
			return ferr
		}
		if errors.As(err, &ke) {
			return fmt.Errorf("%s's host key is not the one recorded in %s: "+
				"if the machine was rebuilt, remove its line there; otherwise something is in the way", host, path)
		}
		return err
	}, nil
}

// sudoersDrop is the grant PrepareSudo writes and RevokeSudo removes.
const sudoersDrop = "/etc/sudoers.d/llmgw-setup"

// sudoGrantLife is how long the grant lasts if nothing removes it sooner.
const sudoGrantLife = time.Hour

// CanSudo reports whether the user can sudo without a password now.
func (g *GoSSH) CanSudo(ctx context.Context) bool {
	_, err := g.Run(ctx, "sudo -n true", nil)
	return err == nil
}

// PrepareSudo makes plain sudo work for the rest of the session: nothing to do
// when it already does, otherwise a sudoers grant written with SudoPassword
// and taken back by RevokeSudo. A grant rather than the password piped into
// every sudo, because the scripts feed their own standard input to the
// commands they run under it.
func (g *GoSSH) PrepareSudo(ctx context.Context) error {
	if g.CanSudo(ctx) {
		return nil
	}
	if g.SudoPassword == "" {
		return fmt.Errorf("%s cannot use sudo without a password, and none was given", g.User)
	}
	script := "sudo -S -p '' sh -c " + shellQuote(sudoGrantScript(g.User))
	if out, err := g.Run(ctx, script, []byte(g.SudoPassword+"\n")); err != nil {
		return fmt.Errorf("sudo with the password given failed on %s: %v: %s", g.Addr, err, strings.TrimSpace(string(out)))
	}
	g.mu.Lock()
	g.sudoGrant = true
	// Not needed again: the grant stands in for it from here on.
	g.SudoPassword = ""
	g.mu.Unlock()
	return nil
}

// RevokeSudo removes the grant PrepareSudo wrote, if it wrote one.
func (g *GoSSH) RevokeSudo(ctx context.Context) error {
	g.mu.Lock()
	granted := g.sudoGrant
	g.mu.Unlock()
	if !granted {
		return nil
	}
	if _, err := g.Run(ctx, "sudo rm -f "+sudoersDrop, nil); err != nil {
		return err
	}
	g.mu.Lock()
	g.sudoGrant = false
	g.mu.Unlock()
	return nil
}

// sudoGrantScript writes the grant, run as root.
func sudoGrantScript(user string) string {
	line := fmt.Sprintf("%s ALL=(ALL) NOPASSWD:ALL", user)
	// Checked by visudo before it is put in place: a broken file in
	// sudoers.d breaks sudo for everyone on the machine.
	//
	// And it removes itself an hour later whatever happens to this process:
	// RevokeSudo is called when the setup finishes, and a setup that crashes
	// or is interrupted never gets there — which would leave a permanent
	// passwordless root behind.
	inner := fmt.Sprintf("umask 337; printf '%%s\\n' %s %s > %[3]s.new && visudo -cf %[3]s.new >/dev/null && mv -f %[3]s.new %[3]s && "+
		"{ systemd-run --quiet --on-active=%[4]d /bin/rm -f %[3]s >/dev/null 2>&1 || "+
		"(setsid nohup sh -c 'sleep %[4]d; rm -f %[3]s' >/dev/null 2>&1 &) ; }",
		shellQuote("# llmgw setup: removed when it finishes, or within the hour"), shellQuote(line), sudoersDrop,
		int(sudoGrantLife.Seconds()))
	return inner
}