package cloud import ( "bufio" "context" "fmt" "io" "os" "path/filepath" "regexp" "strings" "github.com/j3ssie/osmedeus/v5/internal/core" "github.com/j3ssie/osmedeus/v5/internal/runner" "github.com/j3ssie/osmedeus/v5/internal/terminal" "github.com/pkg/sftp" "golang.org/x/crypto/ssh" ) // ansiControlRegex strips non-visual ANSI sequences (cursor, screen, OSC, mode) // while preserving SGR sequences (colors/bold/reset ending in 'm'). // OSC sequences can end with BEL (\x07) or ST (\x1b\\). var ansiControlRegex = regexp.MustCompile(`\x1b\[[0-9;]*[A-HJ-Zadfhijklnqrstu]|\x1b\][^\x1b\x07]*(?:\x07|\x1b\\)|\x1b\(B|\x1b\[\?[0-9;]*[hl]`) // ansiAllRegex strips all ANSI escape sequences including colors (used when color is disabled). var ansiAllRegex = regexp.MustCompile(`\x1b\[[0-9;]*[a-zA-Z]|\x1b\][^\x1b\x07]*(?:\x07|\x1b\\)|\x1b\(B|\x1b\[[\?]?[0-9;]*[hlm]`) // SSHConfig holds the parameters needed to connect to a remote host type SSHConfig struct { Host string Port int User string KeyFile string Password string } // CloudSSHClient wraps a pooled SSH connection for cloud operations. // It provides command execution (blocking and streaming) plus SFTP file transfers. type CloudSSHClient struct { client *ssh.Client poolKey runner.SSHPoolKey config *core.RunnerConfig } // NewCloudSSHClient creates a new SSH client using the global connection pool. func NewCloudSSHClient(ctx context.Context, cfg SSHConfig) (*CloudSSHClient, error) { if cfg.Port == 0 { cfg.Port = 22 } if cfg.User == "" { cfg.User = "root" } runnerCfg := &core.RunnerConfig{ Host: cfg.Host, Port: cfg.Port, User: cfg.User, KeyFile: cfg.KeyFile, Password: cfg.Password, } pool := runner.GetSSHPool() client, poolKey, err := pool.Get(ctx, runnerCfg) if err != nil { return nil, fmt.Errorf("SSH connection to %s:%d failed: %w", cfg.Host, cfg.Port, err) } return &CloudSSHClient{ client: client, poolKey: poolKey, config: runnerCfg, }, nil } // Close releases the connection back to the pool. func (c *CloudSSHClient) Close() { runner.GetSSHPool().Release(c.poolKey) } // RunCommand executes a command on the remote host and returns the output. func (c *CloudSSHClient) RunCommand(ctx context.Context, command string) (string, int, error) { session, err := c.client.NewSession() if err != nil { return "", -1, fmt.Errorf("failed to create session: %w", err) } defer func() { _ = session.Close() }() // Request a PTY so sudo works (non-fatal if unavailable) _ = session.RequestPty("xterm", 80, 200, ssh.TerminalModes{}) output, err := session.CombinedOutput(command) exitCode := 0 if err != nil { if exitErr, ok := err.(*ssh.ExitError); ok { exitCode = exitErr.ExitStatus() } else { return string(output), -1, err } } return string(output), exitCode, nil } // LineCallback is called for each ANSI-stripped output line during streaming. type LineCallback func(line string) // RunCommandStreaming executes a command, streaming stdout/stderr line-by-line with a prefix. // Blocks until the command completes. func (c *CloudSSHClient) RunCommandStreaming(ctx context.Context, command string, prefix string) error { return c.RunCommandStreamingWithCallback(ctx, command, prefix, nil) } // RunCommandStreamingWithCallback is like RunCommandStreaming but calls onLine for each // ANSI-stripped output line, allowing the caller to observe progress without parsing stdout. func (c *CloudSSHClient) RunCommandStreamingWithCallback(ctx context.Context, command string, prefix string, onLine LineCallback) error { session, err := c.client.NewSession() if err != nil { return fmt.Errorf("failed to create session: %w", err) } defer func() { _ = session.Close() }() // Request a PTY so sudo works and output is interleaved (non-fatal if unavailable) _ = session.RequestPty("xterm", 80, 200, ssh.TerminalModes{}) stdout, err := session.StdoutPipe() if err != nil { return fmt.Errorf("failed to get stdout pipe: %w", err) } stderr, err := session.StderrPipe() if err != nil { return fmt.Errorf("failed to get stderr pipe: %w", err) } if err := session.Start(command); err != nil { return fmt.Errorf("failed to start command: %w", err) } prefixStr := terminal.Gray(fmt.Sprintf("[%s] ", prefix)) // oscFragmentRegex catches leftover OSC fragments when the ESC opener was on a previous line // e.g. "11;rgb:1818/1919/2020" from a split OSC 11 background-color response oscFragmentRegex := regexp.MustCompile(`^[0-9]+;[^\x1b]*(?:\x07|\x1b\\)?`) // Stream both pipes with prefix in parallel done := make(chan struct{}, 2) streamLines := func(r io.Reader) { defer func() { done <- struct{}{} }() scanner := bufio.NewScanner(r) scanner.Buffer(make([]byte, 1024*1024), 1024*1024) pendingOSC := false for scanner.Scan() { raw := scanner.Text() // If previous line had an unterminated OSC, this line is the continuation if pendingOSC { pendingOSC = false raw = oscFragmentRegex.ReplaceAllString(raw, "") } // Detect unterminated OSC at end of line (ESC ] without matching terminator) if idx := strings.LastIndex(raw, "\x1b]"); idx >= 0 { tail := raw[idx:] if !strings.Contains(tail, "\x07") && !strings.Contains(tail, "\x1b\\") { raw = raw[:idx] pendingOSC = true } } var line string if terminal.IsColorEnabled() { // Preserve SGR (color) codes; strip only control sequences line = ansiControlRegex.ReplaceAllString(raw, "") } else { line = ansiAllRegex.ReplaceAllString(raw, "") } if line == "" { continue } _, _ = fmt.Fprintf(os.Stdout, "%s%s\n", prefixStr, line) if onLine != nil { onLine(line) } } } go streamLines(stdout) go streamLines(stderr) <-done <-done return session.Wait() } // UploadFile copies a local file to the remote host via SFTP. func (c *CloudSSHClient) UploadFile(localPath, remotePath string) error { sftpClient, err := sftp.NewClient(c.client) if err != nil { return fmt.Errorf("failed to create SFTP client: %w", err) } defer func() { _ = sftpClient.Close() }() // Ensure remote directory exists remoteDir := filepath.Dir(remotePath) _ = sftpClient.MkdirAll(remoteDir) localFile, err := os.Open(localPath) if err != nil { return fmt.Errorf("failed to open local file: %w", err) } defer func() { _ = localFile.Close() }() remoteFile, err := sftpClient.Create(remotePath) if err != nil { return fmt.Errorf("failed to create remote file: %w", err) } defer func() { _ = remoteFile.Close() }() if _, err := io.Copy(remoteFile, localFile); err != nil { return fmt.Errorf("failed to upload file: %w", err) } return nil } // DownloadFile copies a remote file to the local filesystem via SFTP. func (c *CloudSSHClient) DownloadFile(remotePath, localPath string) error { sftpClient, err := sftp.NewClient(c.client) if err != nil { return fmt.Errorf("failed to create SFTP client: %w", err) } defer func() { _ = sftpClient.Close() }() remoteFile, err := sftpClient.Open(remotePath) if err != nil { return fmt.Errorf("failed to open remote file: %w", err) } defer func() { _ = remoteFile.Close() }() // Ensure local directory exists localDir := filepath.Dir(localPath) if err := os.MkdirAll(localDir, 0755); err != nil { return fmt.Errorf("failed to create local directory: %w", err) } localFile, err := os.Create(localPath) if err != nil { return fmt.Errorf("failed to create local file: %w", err) } defer func() { _ = localFile.Close() }() if _, err := io.Copy(localFile, remoteFile); err != nil { return fmt.Errorf("failed to download file: %w", err) } return nil } // DownloadDir recursively downloads a remote directory to local via SFTP. func (c *CloudSSHClient) DownloadDir(remotePath, localPath string) error { sftpClient, err := sftp.NewClient(c.client) if err != nil { return fmt.Errorf("failed to create SFTP client: %w", err) } defer func() { _ = sftpClient.Close() }() return downloadDirRecursive(sftpClient, remotePath, localPath) } func downloadDirRecursive(client *sftp.Client, remotePath, localPath string) error { entries, err := client.ReadDir(remotePath) if err != nil { return fmt.Errorf("failed to list remote dir %s: %w", remotePath, err) } if err := os.MkdirAll(localPath, 0755); err != nil { return fmt.Errorf("failed to create local dir: %w", err) } for _, entry := range entries { remoteEntry := remotePath + "/" + entry.Name() localEntry := filepath.Join(localPath, entry.Name()) if entry.IsDir() { if err := downloadDirRecursive(client, remoteEntry, localEntry); err != nil { return err } } else { remoteFile, err := client.Open(remoteEntry) if err != nil { return fmt.Errorf("failed to open %s: %w", remoteEntry, err) } localFile, err := os.Create(localEntry) if err != nil { _ = remoteFile.Close() return fmt.Errorf("failed to create %s: %w", localEntry, err) } _, copyErr := io.Copy(localFile, remoteFile) _ = remoteFile.Close() _ = localFile.Close() if copyErr != nil { return fmt.Errorf("failed to copy %s: %w", remoteEntry, copyErr) } } } return nil } // ExpandPath expands ~ to the user's home directory func ExpandPath(path string) string { if strings.HasPrefix(path, "~/") { home, err := os.UserHomeDir() if err == nil { return filepath.Join(home, path[2:]) } } return path }