fetch_ml/tests/fixtures/ssh_server.go

// Package tests provides test infrastructure for TUI SSH testing
package tests

import (
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"testing"
	"time"

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

// SSHTestServer manages an SSH test server container for TUI testing
type SSHTestServer struct {
	Host       string
	Port       int
	User       string
	PrivateKey []byte
	Signer     ssh.Signer
	Container  string
	t          *testing.T
}

// NewSSHTestServer creates and starts an SSH test server in a Docker container
func NewSSHTestServer(t *testing.T) *SSHTestServer {
	t.Helper()

	_, filename, _, ok := runtime.Caller(0)
	if !ok {
		t.Fatal("failed to resolve caller path")
	}

	// Navigate to repo root
	repoRoot := filepath.Clean(filepath.Join(filepath.Dir(filename), "..", ".."))

	// Ensure SSH keys exist
	keysDir := filepath.Join(repoRoot, "deployments", "test_keys")
	privateKeyPath := filepath.Join(keysDir, "test_key")

	if _, err := os.Stat(privateKeyPath); os.IsNotExist(err) {
		t.Skipf("SSH test keys not found at %s. Run: ./scripts/testing/gen-ssh-test-keys.sh", keysDir)
	}

	// Read private key
	privateKey, err := os.ReadFile(privateKeyPath)
	if err != nil {
		t.Fatalf("failed to read private key: %v", err)
	}

	signer, err := ssh.ParsePrivateKey(privateKey)
	if err != nil {
		t.Fatalf("failed to parse private key: %v", err)
	}

	server := &SSHTestServer{
		Host:       "localhost",
		Port:       2222,
		User:       "test",
		PrivateKey: privateKey,
		Signer:     signer,
		Container:  "ml-ssh-test",
		t:          t,
	}

	// Verify container is running
	if err := server.waitForContainer(); err != nil {
		t.Skipf("SSH test container not available: %v. Run: docker-compose -f deployments/docker-compose.prod.smoke.yml up -d ssh-test-server", err)
	}

	// Verify SSH connectivity
	if err := server.waitForSSH(); err != nil {
		t.Fatalf("SSH server not reachable: %v", err)
	}

	return server
}

// waitForContainer checks if the SSH container is running
func (s *SSHTestServer) waitForContainer() error {
	cmd := exec.Command("docker", "ps", "-q", "-f", fmt.Sprintf("name=%s", s.Container))
	output, err := cmd.Output()
	if err != nil {
		return fmt.Errorf("failed to check container status: %w", err)
	}
	if len(output) == 0 {
		return fmt.Errorf("container %s is not running", s.Container)
	}
	return nil
}

// waitForSSH waits for the SSH server to accept connections
func (s *SSHTestServer) waitForSSH() error {
	config := &ssh.ClientConfig{
		User: s.User,
		Auth: []ssh.AuthMethod{
			ssh.PublicKeys(s.Signer),
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(), //nolint:gosec // Test only
		Timeout:         5 * time.Second,
	}

	addr := fmt.Sprintf("%s:%d", s.Host, s.Port)

	// Retry with backoff
	for i := 0; i < 10; i++ {
		client, err := ssh.Dial("tcp", addr, config)
		if err == nil {
			client.Close()
			return nil
		}
		time.Sleep(500 * time.Millisecond)
	}

	return fmt.Errorf("SSH server not available after retries")
}

// NewClient creates a new SSH client connection
func (s *SSHTestServer) NewClient() (*ssh.Client, error) {
	config := &ssh.ClientConfig{
		User: s.User,
		Auth: []ssh.AuthMethod{
			ssh.PublicKeys(s.Signer),
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(), //nolint:gosec // Test only
		Timeout:         10 * time.Second,
	}

	addr := fmt.Sprintf("%s:%d", s.Host, s.Port)
	return ssh.Dial("tcp", addr, config)
}

// Stop stops the SSH test server container
func (s *SSHTestServer) Stop() {
	// Container is managed by docker-compose, so we don't stop it here
	// This allows multiple tests to reuse the same container
}

// Cleanup stops and removes the SSH test container
func (s *SSHTestServer) Cleanup() {
	cmd := exec.Command("docker", "stop", s.Container)
	_ = cmd.Run()

	cmd = exec.Command("docker", "rm", s.Container)
	_ = cmd.Run()
}

// Exec runs a command on the SSH server and returns output
func (s *SSHTestServer) Exec(cmd string) (string, error) {
	client, err := s.NewClient()
	if err != nil {
		return "", fmt.Errorf("failed to connect: %w", err)
	}
	defer client.Close()

	session, err := client.NewSession()
	if err != nil {
		return "", fmt.Errorf("failed to create session: %w", err)
	}
	defer session.Close()

	output, err := session.CombinedOutput(cmd)
	if err != nil {
		return string(output), fmt.Errorf("command failed: %w", err)
	}

	return string(output), nil
}

// ExecWithPTY runs a command with PTY allocation for interactive TUI testing
func (s *SSHTestServer) ExecWithPTY(cmd string, term string, width, height int) (*ssh.Session, error) {
	client, err := s.NewClient()
	if err != nil {
		return nil, fmt.Errorf("failed to connect: %w", err)
	}

	session, err := client.NewSession()
	if err != nil {
		client.Close()
		return nil, fmt.Errorf("failed to create session: %w", err)
	}

	// Request PTY
	modes := ssh.TerminalModes{
		ssh.ECHO:          0,
		ssh.TTY_OP_ISPEED: 14400,
		ssh.TTY_OP_OSPEED: 14400,
	}

	if err := session.RequestPty(term, width, height, modes); err != nil {
		session.Close()
		client.Close()
		return nil, fmt.Errorf("failed to request pty: %w", err)
	}

	if err := session.Start(cmd); err != nil {
		session.Close()
		client.Close()
		return nil, fmt.Errorf("failed to start command: %w", err)
	}

	return session, nil
}