package runner import ( "context" "fmt" "io" "os" "os/exec" "path/filepath" "strings" "time" "github.com/j3ssie/osmedeus/v5/internal/core" "github.com/pkg/sftp" "golang.org/x/crypto/ssh" ) // SSHRunner executes commands on a remote machine via SSH type SSHRunner struct { config *core.RunnerConfig binaryPath string client *ssh.Client remoteDir string pooled bool // true if connection is from pool poolKey SSHPoolKey // key for releasing back to pool } // NewSSHRunner creates a new SSH runner func NewSSHRunner(config *core.RunnerConfig, binaryPath string) (*SSHRunner, error) { if config.Host == "" { return nil, fmt.Errorf("SSH runner requires host to be specified") } if config.User == "" { return nil, fmt.Errorf("SSH runner requires user to be specified") } return &SSHRunner{ config: config, binaryPath: binaryPath, remoteDir: "~/.osmedeus-remote", }, nil } // Setup establishes SSH connection and copies the binary // Uses connection pooling to reuse SSH connections across steps func (r *SSHRunner) Setup(ctx context.Context) error { // Get connection from pool pool := GetSSHPool() client, poolKey, err := pool.Get(ctx, r.config) if err != nil { return err } r.client = client r.pooled = true r.poolKey = poolKey // Create remote directory if _, err := r.runCommand(ctx, "mkdir -p "+r.remoteDir); err != nil { return fmt.Errorf("failed to create remote directory: %w", err) } // Copy binary to remote machine if it exists if r.binaryPath != "" { if err := r.copyBinary(); err != nil { // Non-fatal, log and continue fmt.Printf("Warning: failed to copy binary to remote: %v\n", err) } } return nil } // expandPath expands ~ to home directory func (r *SSHRunner) expandPath(path string) string { if strings.HasPrefix(path, "~/") { home, err := os.UserHomeDir() if err == nil { return filepath.Join(home, path[2:]) } } return path } // copyBinary copies the osmedeus binary to the remote machine using SFTP func (r *SSHRunner) copyBinary() error { if r.client == nil { return fmt.Errorf("SSH client not connected") } // Create SFTP client sftpClient, err := sftp.NewClient(r.client) if err != nil { return fmt.Errorf("failed to create SFTP client: %w", err) } defer func() { _ = sftpClient.Close() }() // Expand remote directory homeDir, err := r.getRemoteHome() if err != nil { return fmt.Errorf("failed to get remote home: %w", err) } remotePath := strings.Replace(r.remoteDir, "~", homeDir, 1) // Create remote directory _ = sftpClient.MkdirAll(remotePath) // Open local file localFile, err := os.Open(r.binaryPath) if err != nil { return fmt.Errorf("failed to open local binary: %w", err) } defer func() { _ = localFile.Close() }() // Get local file info for size localInfo, err := localFile.Stat() if err != nil { return fmt.Errorf("failed to stat local binary: %w", err) } // Create remote file destPath := filepath.Join(remotePath, "osmedeus") remoteFile, err := sftpClient.Create(destPath) if err != nil { return fmt.Errorf("failed to create remote file: %w", err) } defer func() { _ = remoteFile.Close() }() // Copy file _, err = io.Copy(remoteFile, localFile) if err != nil { return fmt.Errorf("failed to copy binary: %w", err) } // Make executable if err := sftpClient.Chmod(destPath, 0755); err != nil { return fmt.Errorf("failed to chmod binary: %w", err) } fmt.Printf("Copied binary to remote (%d bytes)\n", localInfo.Size()) return nil } // getRemoteHome gets the home directory on the remote machine func (r *SSHRunner) getRemoteHome() (string, error) { // Use background context for internal utility calls output, err := r.runCommand(context.Background(), "echo $HOME") if err != nil { return "", err } return strings.TrimSpace(output), nil } // runCommand executes a command on the remote machine with context support func (r *SSHRunner) runCommand(ctx context.Context, command string) (string, error) { if r.client == nil { return "", fmt.Errorf("SSH client not connected") } session, err := r.client.NewSession() if err != nil { return "", fmt.Errorf("failed to create SSH session: %w", err) } defer func() { _ = session.Close() }() // Set working directory if specified if r.config.WorkDir != "" { command = fmt.Sprintf("cd %s && %s", r.config.WorkDir, command) } // Enforce context deadline using timeout command on remote host if deadline, ok := ctx.Deadline(); ok { remaining := time.Until(deadline) if remaining > 0 { secs := int(remaining.Seconds()) if secs > 0 { command = fmt.Sprintf("timeout %ds %s", secs, command) } } } output, err := session.CombinedOutput(command) return string(output), err } // Execute runs a command on the remote machine func (r *SSHRunner) Execute(ctx context.Context, command string) (*CommandResult, error) { output, err := r.runCommand(ctx, command) exitCode := 0 if err != nil { if exitErr, ok := err.(*ssh.ExitError); ok { exitCode = exitErr.ExitStatus() } } return &CommandResult{ Output: output, ExitCode: exitCode, Error: err, }, nil } // Cleanup releases the SSH connection back to the pool func (r *SSHRunner) Cleanup(ctx context.Context) error { if r.client != nil { if r.pooled { // Release back to pool for reuse GetSSHPool().Release(r.poolKey) } else { // Direct close for non-pooled connections _ = r.client.Close() } r.client = nil } return nil } // getPort returns the SSH port, defaulting to 22 func (r *SSHRunner) getPort() int { if r.config.Port == 0 { return 22 } return r.config.Port } // CopyFromRemote copies a file from the SSH host to local using rsync func (r *SSHRunner) CopyFromRemote(ctx context.Context, remotePath, localPath string) error { // Ensure parent directory exists if dir := filepath.Dir(localPath); dir != "" && dir != "." { if err := os.MkdirAll(dir, 0755); err != nil { return fmt.Errorf("failed to create directory: %w", err) } } // Build rsync source: user@host:path source := fmt.Sprintf("%s@%s:%s", r.config.User, r.config.Host, remotePath) // Build rsync command with SSH options args := []string{"-avz", "-e"} if r.config.KeyFile != "" { keyPath := r.expandPath(r.config.KeyFile) args = append(args, fmt.Sprintf("ssh -i %s -p %d -o StrictHostKeyChecking=no", keyPath, r.getPort())) } else { args = append(args, fmt.Sprintf("ssh -p %d -o StrictHostKeyChecking=no", r.getPort())) } args = append(args, source, localPath) cmd := exec.CommandContext(ctx, "rsync", args...) output, err := cmd.CombinedOutput() if err != nil { return fmt.Errorf("rsync failed: %w, output: %s", err, string(output)) } return nil } // Type returns the runner type func (r *SSHRunner) Type() core.RunnerType { return core.RunnerTypeSSH } // IsRemote returns true since SSH runner runs on a remote machine func (r *SSHRunner) IsRemote() bool { return true } // GetRemoteDir returns the remote directory where binary is stored func (r *SSHRunner) GetRemoteDir() string { return r.remoteDir }