orca/internal/proxmox/ssh_session_test.go

package proxmox

import (
	"bytes"
	"context"
	"crypto/ed25519"
	"crypto/rand"
	"errors"
	"log/slog"
	"net"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"testing"
	"time"

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

type fakeSSHServer struct {
	listener net.Listener
	config   *ssh.ServerConfig
	done     chan struct{}

	mu                  sync.Mutex
	state               map[string]string
	authDir             string
	forceSudoersInvalid bool
}

func newFakeSSHServer(t *testing.T) *fakeSSHServer {
	t.Helper()
	_, priv, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		t.Fatalf("ed25519 gen: %v", err)
	}
	hostSigner, err := ssh.NewSignerFromKey(priv)
	if err != nil {
		t.Fatalf("ssh signer: %v", err)
	}
	config := &ssh.ServerConfig{
		PasswordCallback: func(c ssh.ConnMetadata, password []byte) (*ssh.Permissions, error) {
			if string(password) != "pw" {
				return nil, errors.New("invalid password")
			}
			return nil, nil
		},
	}
	config.AddHostKey(hostSigner)

	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen: %v", err)
	}
	srv := &fakeSSHServer{
		listener: ln,
		config:   config,
		done:     make(chan struct{}),
		state:    make(map[string]string),
		authDir:  t.TempDir(),
	}
	go srv.serve()
	return srv
}

func (s *fakeSSHServer) addr() string { return s.listener.Addr().String() }

func (s *fakeSSHServer) serve() {
	for {
		conn, err := s.listener.Accept()
		if err != nil {
			close(s.done)
			return
		}
		go s.handle(conn)
	}
}

func (s *fakeSSHServer) handle(netConn net.Conn) {
	defer netConn.Close()
	_, chans, reqs, err := ssh.NewServerConn(netConn, s.config)
	if err != nil {
		return
	}
	go ssh.DiscardRequests(reqs)
	for newChan := range chans {
		if newChan.ChannelType() != "session" {
			newChan.Reject(ssh.UnknownChannelType, "only session")
			continue
		}
		go s.handleSession(newChan)
	}
}

func (s *fakeSSHServer) handleSession(newChan ssh.NewChannel) {
	ch, reqs, err := newChan.Accept()
	if err != nil {
		return
	}
	defer ch.Close()
	for req := range reqs {
		switch req.Type {
		case "exec":
			var execReq struct{ Command string }
			if err := ssh.Unmarshal(req.Payload, &execReq); err != nil {
				req.Reply(false, nil)
				continue
			}
			req.Reply(true, nil)
			out, code := s.runCommand(execReq.Command)
			_, _ = ch.Write(out)
			_, _ = ch.SendRequest("exit-status", false, ssh.Marshal(struct{ Code uint32 }{uint32(code)}))
			_ = ch.Close()
		default:
			req.Reply(false, nil)
		}
	}
}

func (s *fakeSSHServer) runCommand(cmd string) ([]byte, int) {
	s.mu.Lock()
	defer s.mu.Unlock()
	trimmed := strings.TrimSpace(cmd)
	switch {
	case trimmed == "echo hello":
		return []byte("hello\n"), 0
	case strings.HasPrefix(trimmed, "exit "):
		return nil, 1
	case strings.HasPrefix(trimmed, "id -u "):
		return []byte("1000\n"), 0
	case strings.Contains(trimmed, "pveum role list") || strings.Contains(trimmed, "pveum role add"):
		s.state["pve_role:"+extractField(trimmed, "add ", " ")] = "ok"
		return nil, 0
	case strings.Contains(trimmed, "pveum user list") || strings.Contains(trimmed, "pveum user add"):
		s.state["pve_user:orca@pam"] = "ok"
		return nil, 0
	case strings.Contains(trimmed, "pveum acl modify"):
		s.state["pve_acl"] = "ok"
		return nil, 0
	case strings.HasPrefix(trimmed, "mkdir -p ") && strings.Contains(trimmed, "authorized_keys"):
		return s.handleAuthKeyDeploy(trimmed)
	case strings.HasPrefix(trimmed, "cat > /etc/sudoers.d/"):
		return s.handleSudoersWrite(trimmed), 0
	case strings.HasPrefix(trimmed, "visudo -cf /etc/sudoers.d/orca"):
		force := s.forceSudoersInvalid
		if force || s.state["sudoers_valid"] != "true" {
			return []byte("/etc/sudoers.d/orca: syntax error\n"), 1
		}
		return []byte("/etc/sudoers.d/orca: parsed OK\n"), 0
	case strings.HasPrefix(trimmed, "cat /") && strings.HasSuffix(trimmed, "/authorized_keys"):
		return s.readAuthFile(trimmed[4:]), 0
	case strings.HasPrefix(trimmed, "cat /") && strings.Contains(trimmed, "/orca"):
		return s.readSudoers(trimmed[4:]), 0
	default:
		return []byte("sh: command not found\n"), 127
	}
}

func (s *fakeSSHServer) handleAuthKeyDeploy(cmd string) ([]byte, int) {
	parts := strings.Split(cmd, "'")
	var pubLine string
	if len(parts) >= 2 {
		pubLine = parts[1]
	}
	authPath := filepath.Join(s.authDir, "authorized_keys")
	existing := string(s.readFile(authPath))
	if !strings.Contains(existing, pubLine) {
		existing += pubLine + "\n"
	}
	if err := os.WriteFile(authPath, []byte(existing), 0o600); err != nil {
		return []byte("mkdir: permission denied\n"), 1
	}
	return nil, 0
}

func (s *fakeSSHServer) readAuthFile(path string) []byte {
	if strings.HasSuffix(path, "/authorized_keys") {
		return s.readFile(filepath.Join(s.authDir, "authorized_keys"))
	}
	return []byte("cat: " + path + ": No such file or directory\n")
}

func (s *fakeSSHServer) handleSudoersWrite(cmd string) []byte {
	idx := strings.Index(cmd, "\n")
	if idx < 0 {
		return []byte("sh: bad heredoc\n")
	}
	content := cmd[idx+1:]
	if end := strings.Index(content, "ORCA_SUDOERS_EOF"); end >= 0 {
		content = content[:end]
	}
	s.state["sudoers_content"] = content
	s.state["sudoers_valid"] = "true"
	return nil
}

func (s *fakeSSHServer) readSudoers(path string) []byte {
	if v, ok := s.state["sudoers_content"]; ok {
		return []byte(v)
	}
	return []byte("cat: " + path + ": No such file or directory\n")
}

func (s *fakeSSHServer) readFile(path string) []byte {
	b, _ := os.ReadFile(path)
	return b
}

func (s *fakeSSHServer) close() {
	s.listener.Close()
	<-s.done
}

func extractField(s, after, until string) string {
	i := strings.Index(s, after)
	if i < 0 {
		return ""
	}
	rest := s[i+len(after):]
	j := strings.Index(rest, until)
	if j < 0 {
		return rest
	}
	return rest[:j]
}

func fakeSSHClient(t *testing.T, srv *fakeSSHServer) *ssh.Client {
	t.Helper()
	config := &ssh.ClientConfig{
		User:            "root",
		Auth:            []ssh.AuthMethod{ssh.Password("pw")},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		Timeout:         5 * time.Second,
	}
	client, err := ssh.Dial("tcp", srv.addr(), config)
	if err != nil {
		t.Fatalf("ssh.Dial: %v", err)
	}
	return client
}

func withSessionRunner(t *testing.T, conn *ssh.Client) {
	t.Helper()
	orig := sessionRunner
	t.Cleanup(func() { sessionRunner = orig })
	sessionRunner = &sshSessionRunner{client: conn}
}

func TestRunRemote_Success(t *testing.T) {
	srv := newFakeSSHServer(t)
	defer srv.close()
	conn := fakeSSHClient(t, srv)
	defer conn.Close()
	withSessionRunner(t, conn)
	out, err := runRemote("echo hello")
	if err != nil {
		t.Fatalf("runRemote: %v", err)
	}
	if strings.TrimSpace(string(out)) != "hello" {
		t.Errorf("output = %q, want hello", strings.TrimSpace(string(out)))
	}
}

func TestRunRemote_Failure(t *testing.T) {
	srv := newFakeSSHServer(t)
	defer srv.close()
	conn := fakeSSHClient(t, srv)
	defer conn.Close()
	withSessionRunner(t, conn)
	_, err := runRemote("exit 7")
	if err == nil {
		t.Fatal("expected error for non-zero exit")
	}
	if !strings.Contains(err.Error(), "run") {
		t.Errorf("error should mention run, got: %v", err)
	}
}

func TestDeployPubKey_Success(t *testing.T) {
	srv := newFakeSSHServer(t)
	defer srv.close()
	conn := fakeSSHClient(t, srv)
	defer conn.Close()
	withSessionRunner(t, conn)

	if err := deployPubKey("orca", "ssh-ed25519 AAAA test@orca"); err != nil {
		t.Fatalf("deployPubKey: %v", err)
	}
	out := srv.readFile(filepath.Join(srv.authDir, "authorized_keys"))
	if !strings.Contains(string(out), "ssh-ed25519 AAAA test@orca") {
		t.Errorf("auth file does not contain the key: %s", out)
	}
}

func TestDeployPubKey_Idempotent(t *testing.T) {
	srv := newFakeSSHServer(t)
	defer srv.close()
	conn := fakeSSHClient(t, srv)
	defer conn.Close()
	withSessionRunner(t, conn)

	if err := deployPubKey("orca", "ssh-ed25519 AAAA test@orca"); err != nil {
		t.Fatalf("first deploy: %v", err)
	}
	if err := deployPubKey("orca", "ssh-ed25519 AAAA test@orca"); err != nil {
		t.Fatalf("second deploy: %v", err)
	}
	out := srv.readFile(filepath.Join(srv.authDir, "authorized_keys"))
	if cnt := strings.Count(string(out), "ssh-ed25519 AAAA test@orca"); cnt != 1 {
		t.Errorf("key count = %d, want 1 (idempotent)", cnt)
	}
}

func TestCreateLinuxUser_Success(t *testing.T) {
	srv := newFakeSSHServer(t)
	defer srv.close()
	conn := fakeSSHClient(t, srv)
	defer conn.Close()
	withSessionRunner(t, conn)
	if err := createLinuxUser("orca"); err != nil {
		t.Fatalf("createLinuxUser: %v", err)
	}
}

func TestCreatePVERole_Success(t *testing.T) {
	srv := newFakeSSHServer(t)
	defer srv.close()
	conn := fakeSSHClient(t, srv)
	defer conn.Close()
	withSessionRunner(t, conn)
	if err := createPVERole("OrcaOperator"); err != nil {
		t.Fatalf("createPVERole: %v", err)
	}
}

func TestCreatePVEUser_Success(t *testing.T) {
	srv := newFakeSSHServer(t)
	defer srv.close()
	conn := fakeSSHClient(t, srv)
	defer conn.Close()
	withSessionRunner(t, conn)
	if err := createPVEUser("orca"); err != nil {
		t.Fatalf("createPVEUser: %v", err)
	}
}

func TestAssignPVEACL_Success(t *testing.T) {
	srv := newFakeSSHServer(t)
	defer srv.close()
	conn := fakeSSHClient(t, srv)
	defer conn.Close()
	withSessionRunner(t, conn)
	if err := assignPVEACL("orca", "OrcaOperator"); err != nil {
		t.Fatalf("assignPVEACL: %v", err)
	}
}

func TestWriteSudoers_Success(t *testing.T) {
	srv := newFakeSSHServer(t)
	defer srv.close()
	conn := fakeSSHClient(t, srv)
	defer conn.Close()
	withSessionRunner(t, conn)

	if err := writeSudoers("orca"); err != nil {
		t.Fatalf("writeSudoers: %v", err)
	}
	if srv.state["sudoers_valid"] != "true" {
		t.Error("sudoers not marked valid")
	}
	if !strings.Contains(srv.state["sudoers_content"], "orca ALL=(root) NOPASSWD: NOEXEC: /usr/bin/pct") {
		t.Errorf("sudoers content missing pct: %s", srv.state["sudoers_content"])
	}
}

func TestValidateSudoers_ParsedOK(t *testing.T) {
	srv := newFakeSSHServer(t)
	defer srv.close()
	conn := fakeSSHClient(t, srv)
	defer conn.Close()
	withSessionRunner(t, conn)

	srv.state["sudoers_valid"] = "true"
	if err := validateSudoers(); err != nil {
		t.Errorf("validateSudoers: %v", err)
	}
}

func TestValidateSudoers_Failure(t *testing.T) {
	srv := newFakeSSHServer(t)
	defer srv.close()
	conn := fakeSSHClient(t, srv)
	defer conn.Close()
	withSessionRunner(t, conn)

	srv.state["sudoers_valid"] = "false"
	if err := validateSudoers(); err == nil {
		t.Error("expected error for invalid sudoers")
	}
}

type staticDialer struct {
	client *ssh.Client
}

func (d *staticDialer) DialContext(ctx context.Context, network, addr string, config *ssh.ClientConfig) (*ssh.Client, error) {
	return d.client, nil
}

type funcDialer struct {
	fn func(ctx context.Context, network, addr string, config *ssh.ClientConfig) (*ssh.Client, error)
}

func (d *funcDialer) DialContext(ctx context.Context, network, addr string, config *ssh.ClientConfig) (*ssh.Client, error) {
	return d.fn(ctx, network, addr, config)
}

func TestBootstrapProxmox_FullFlow_Success(t *testing.T) {
	srv := newFakeSSHServer(t)
	defer srv.close()

	home := t.TempDir()
	t.Setenv("ORCA_HOME", home)
	if err := os.WriteFile(filepath.Join(home, "known_hosts"), []byte{}, 0o600); err != nil {
		t.Fatalf("create known_hosts: %v", err)
	}

	orig := sshDialer
	defer func() { sshDialer = orig }()
	origRunner := sessionRunner
	defer func() { sessionRunner = origRunner }()
	sessionRunner = nil
	sshDialer = &staticDialer{client: fakeSSHClient(t, srv)}

	host, _, _ := net.SplitHostPort(srv.addr())

	var logBuf bytes.Buffer
	result, err := BootstrapProxmox(t.Context(), Options{
		Host:     host,
		Password: "pw",
		Logger:   slog.New(slog.NewTextHandler(&logBuf, nil)),
	})
	if err != nil {
		t.Fatalf("BootstrapProxmox: %v", err)
	}
	if result == nil {
		t.Fatal("result is nil")
	}
	if result.NodeName != host {
		t.Errorf("NodeName = %q, want %q", result.NodeName, host)
	}
	if result.NodeAddress != host+":8443" {
		t.Errorf("NodeAddress = %q, want %q:8443", result.NodeAddress, host)
	}
	if !strings.Contains(logBuf.String(), "proxmox.bootstrap_ok") {
		t.Errorf("expected bootstrap_ok log, got: %s", logBuf.String())
	}
}

func TestBootstrapProxmox_FullFlow_DeployPubKeyFails(t *testing.T) {
	srv := newFakeSSHServer(t)
	defer srv.close()

	home := t.TempDir()
	t.Setenv("ORCA_HOME", home)
	if err := os.WriteFile(filepath.Join(home, "known_hosts"), []byte{}, 0o600); err != nil {
		t.Fatalf("create known_hosts: %v", err)
	}

	orig := sshDialer
	defer func() { sshDialer = orig }()
	origRunner := sessionRunner
	defer func() { sessionRunner = origRunner }()
	sessionRunner = nil

	// Use a real client that connects to a server which will reject deploy
	// by returning a non-zero exit for the mkdir command. We achieve this
	// by using a dialer that returns a client to a server whose authDir
	// is read-only — but simpler: just use a fresh server that errors on
	// authorized_keys commands via a custom server. We reuse newFakeSSHServer
	// but sabotage it by pointing authDir to a read-only location.
	conn := fakeSSHClient(t, srv)
	defer conn.Close()
	sshDialer = &staticDialer{client: conn}

	host, _, _ := net.SplitHostPort(srv.addr())

	// Make authDir unwritable so deployPubKey's mkdir handler fails.
	srv.authDir = "/proc/1/forbidden-orca-test"

	_, err := BootstrapProxmox(t.Context(), Options{
		Host:     host,
		Password: "pw",
	})
	if err == nil {
		t.Fatal("expected error from deployPubKey failure")
	}
	if !strings.Contains(err.Error(), "deploy pubkey") {
		t.Errorf("error should mention deploy pubkey, got: %v", err)
	}
}