ssh.go

package main

import (
	"context"
	"fmt"
	"net"
	"os"
	"time"

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

// SSHExecutor runs commands on a remote host via SSH.
type SSHExecutor struct {
	client *ssh.Client
}

func NewSSHExecutor(host, user, keyFile, password string, port int) (*SSHExecutor, error) {
	if port == 0 {
		port = 22
	}

	var authMethods []ssh.AuthMethod

	// Key-based auth
	if keyFile != "" {
		keyData, err := os.ReadFile(keyFile)
		if err != nil {
			return nil, fmt.Errorf("reading SSH key: %w", err)
		}
		signer, err := ssh.ParsePrivateKey(keyData)
		if err != nil {
			return nil, fmt.Errorf("parsing SSH key: %w", err)
		}
		authMethods = append(authMethods, ssh.PublicKeys(signer))
	}

	// Password auth
	if password != "" {
		authMethods = append(authMethods, ssh.Password(password))
	}

	if len(authMethods) == 0 {
		return nil, fmt.Errorf("no SSH auth method provided: set key_file or password")
	}

	config := &ssh.ClientConfig{
		User:            user,
		Auth:            authMethods,
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		Timeout:         10 * time.Second,
	}

	addr := net.JoinHostPort(host, fmt.Sprintf("%d", port))
	client, err := ssh.Dial("tcp", addr, config)
	if err != nil {
		return nil, fmt.Errorf("SSH connection to %s failed: %w", addr, err)
	}

	return &SSHExecutor{client: client}, nil
}

func (e *SSHExecutor) Exec(command string, timeout int) (string, error) {
	session, err := e.client.NewSession()
	if err != nil {
		return "", fmt.Errorf("creating SSH session: %w", err)
	}
	defer session.Close()

	if timeout > 0 {
		ctx, cancel := context.WithTimeout(context.Background(), time.Duration(timeout)*time.Second)
		defer cancel()

		done := make(chan struct{})
		var output []byte
		var runErr error

		go func() {
			output, runErr = session.CombinedOutput(command)
			close(done)
		}()

		select {
		case <-done:
			if runErr != nil {
				return string(output) + "\nError: " + runErr.Error(), nil
			}
			return string(output), nil
		case <-ctx.Done():
			session.Signal(ssh.SIGKILL)
			return "", fmt.Errorf("command timed out after %d seconds", timeout)
		}
	}

	output, err := session.CombinedOutput(command)
	if err != nil {
		return string(output) + "\nError: " + err.Error(), nil
	}
	return string(output), nil
}

func (e *SSHExecutor) Close() error {
	if e.client != nil {
		return e.client.Close()
	}
	return nil
}