aiops/backend/pkg/remoteexec/executor.go

// Package remoteexec 提供 Go 原生 SSH 远程命令执行和文件上传能力。
// 基于 golang.org/x/crypto/ssh，零外部依赖，替代 ansible-playbook。
// 支持密码认证、SSH 密钥认证、TOFU 主机密钥验证。
package remoteexec

import (
	"bytes"
	"context"
	"fmt"
	"log"
	"net"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

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

// ============================================
// 类型定义
// ============================================

// AuthMethod SSH 认证方式
type AuthMethod int

const (
	AuthPassword AuthMethod = iota // 密码认证
	AuthKey                        // SSH 私钥认证
)

// SSHConfig SSH 连接配置
type SSHConfig struct {
	Host     string        // 目标 IP 或域名
	Port     int           // SSH 端口（默认 22）
	User     string        // SSH 用户（默认 root）
	AuthType AuthMethod    // 认证方式
	Password string        // 密码（AuthPassword 时使用）
	KeyData  string        // SSH 私钥内容（AuthKey 时使用）
	Timeout  time.Duration // 连接超时（默认 10s）
}

// ExecResult 远程命令执行结果
type ExecResult struct {
	Host     string `json:"host"`     // 目标主机
	Port     int    `json:"port"`     // 目标端口
	Output   string `json:"output"`   // stdout+stderr 合并输出
	ExitCode int    `json:"exit_code"` // 退出码（0=成功）
	Duration int64  `json:"duration"` // 执行耗时（毫秒）
	Error    string `json:"error,omitempty"` // 错误信息（连接失败等）
	Success  bool   `json:"success"`  // 是否成功
}

// UploadResult 文件上传结果
type UploadResult struct {
	Host     string `json:"host"`
	Port     int    `json:"port"`
	LocalPath  string `json:"local_path"`
	RemotePath string `json:"remote_path"`
	Success    bool   `json:"success"`
	Error      string `json:"error,omitempty"`
}

// ============================================
// 核心函数
// ============================================

// Exec 在单台远程主机上执行命令。
// 返回合并的 stdout+stderr 输出、退出码和错误。
func Exec(cfg SSHConfig, command string, timeout time.Duration) (*ExecResult, error) {
	start := time.Now()
	result := &ExecResult{
		Host: cfg.Host,
		Port: cfg.Port,
	}

	client, err := connect(cfg)
	if err != nil {
		result.Duration = time.Since(start).Milliseconds()
		result.Error = err.Error()
		return result, err
	}
	defer client.Close()

	// 创建 session
	session, err := client.NewSession()
	if err != nil {
		result.Duration = time.Since(start).Milliseconds()
		result.Error = fmt.Sprintf("创建SSH会话失败: %v", err)
		return result, err
	}
	defer session.Close()

	// 设置超时
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()

	// 在 goroutine 中执行，支持超时
	type execOut struct {
		output []byte
		err    error
	}
	ch := make(chan execOut, 1)

	go func() {
		var stdout, stderr bytes.Buffer
		session.Stdout = &stdout
		session.Stderr = &stderr
		err := session.Run(command)

		// 合并输出
		output := stdout.Bytes()
		if stderr.Len() > 0 {
			if len(output) > 0 {
				output = append(output, '\n')
			}
			output = append(output, stderr.Bytes()...)
		}
		ch <- execOut{output: output, err: err}
	}()

	select {
	case <-ctx.Done():
		result.Duration = time.Since(start).Milliseconds()
		result.Error = fmt.Sprintf("执行超时（%v）", timeout)
		result.ExitCode = -1
		return result, fmt.Errorf("exec timeout after %v", timeout)
	case out := <-ch:
		result.Duration = time.Since(start).Milliseconds()
		result.Output = string(out.output)

		if out.err != nil {
			if exitErr, ok := out.err.(*ssh.ExitError); ok {
				result.ExitCode = exitErr.ExitStatus()
				result.Success = false
			} else {
				result.ExitCode = -1
				result.Error = out.err.Error()
				result.Success = false
				return result, out.err
			}
		} else {
			result.ExitCode = 0
			result.Success = true
		}
	}

	return result, nil
}

// BatchExec 在多台远程主机上并发执行相同命令。
// maxConcurrency 控制最大并发数（0 = 无限制）。
func BatchExec(servers []SSHConfig, command string, timeout time.Duration, maxConcurrency int) []ExecResult {
	if maxConcurrency <= 0 {
		maxConcurrency = len(servers)
	}

	results := make([]ExecResult, len(servers))
	sem := make(chan struct{}, maxConcurrency)
	var wg sync.WaitGroup

	for i, cfg := range servers {
		wg.Add(1)
		go func(idx int, config SSHConfig) {
			defer wg.Done()
			sem <- struct{}{}
			defer func() { <-sem }()

			result, err := Exec(config, command, timeout)
			if err != nil && result == nil {
				results[idx] = ExecResult{
					Host:    config.Host,
					Port:    config.Port,
					Error:   err.Error(),
					Success: false,
				}
			} else {
				results[idx] = *result
			}
		}(i, cfg)
	}
	wg.Wait()

	return results
}

// UploadFile 上传本地文件到远程主机（使用 SFTP/SCP）。
// 采用标准 SSH exec + base64 编码方式，避免依赖 SFTP 子协议。
// 对于大文件（>1MB），建议改用 sftp 方式。
func UploadFile(cfg SSHConfig, localPath, remotePath string) (*UploadResult, error) {
	result := &UploadResult{
		Host:       cfg.Host,
		Port:       cfg.Port,
		LocalPath:  localPath,
		RemotePath: remotePath,
	}

	// 读取本地文件
	data, err := os.ReadFile(localPath)
	if err != nil {
		result.Error = fmt.Sprintf("读取本地文件失败: %v", err)
		return result, err
	}

	if len(data) > 10*1024*1024 { // 10MB 限制
		result.Error = "文件超过 10MB，请使用其他传输方式"
		return result, fmt.Errorf("file too large: %d bytes", len(data))
	}

	// 通过 SSH exec + base64 方式传输（兼容性好，不需要 SFTP）
	client, err := connect(cfg)
	if err != nil {
		result.Error = err.Error()
		return result, err
	}
	defer client.Close()

	// 步骤1: 确保远程目录存在
	remoteDir := filepath.Dir(remotePath)
	mkdirCmd := fmt.Sprintf("mkdir -p %s", remoteDir)
	Exec(cfg, mkdirCmd, 30*time.Second)

	// 步骤2: base64 编码文件后写入
	encodedLen := len(data) // 实际要分块，这里简化处理
	_ = encodedLen

	session, err := client.NewSession()
	if err != nil {
		result.Error = fmt.Sprintf("创建SSH会话失败: %v", err)
		return result, err
	}
	defer session.Close()

	// 将文件内容直接写入远程（通过 stdin）
	session.Stdin = bytes.NewReader(data)
	writeCmd := fmt.Sprintf("cat > %s", remotePath)

	var stderr bytes.Buffer
	session.Stderr = &stderr

	if err := session.Run(writeCmd); err != nil {
		result.Error = fmt.Sprintf("写入远程文件失败: %v, stderr: %s", err, stderr.String())
		return result, err
	}

	// 验证文件已写入
	verifySession, _ := client.NewSession()
	verifyOut, verifyErr := verifySession.Output(fmt.Sprintf("wc -c < %s", remotePath))
	verifySession.Close()

	if verifyErr == nil {
		log.Printf("[RemoteExec] Uploaded %s → %s:%s (%s bytes written)", 
			localPath, cfg.Host, remotePath, strings.TrimSpace(string(verifyOut)))
	}

	result.Success = true
	return result, nil
}

// BatchUpload 批量上传文件到多台服务器（并发）。
func BatchUpload(servers []SSHConfig, localPath, remotePath string, maxConcurrency int) []UploadResult {
	if maxConcurrency <= 0 {
		maxConcurrency = len(servers)
	}

	results := make([]UploadResult, len(servers))
	sem := make(chan struct{}, maxConcurrency)
	var wg sync.WaitGroup

	for i, cfg := range servers {
		wg.Add(1)
		go func(idx int, config SSHConfig) {
			defer wg.Done()
			sem <- struct{}{}
			defer func() { <-sem }()

			result, err := UploadFile(config, localPath, remotePath)
			if err != nil && result == nil {
				results[idx] = UploadResult{
					Host:       config.Host,
					Port:       config.Port,
					LocalPath:  localPath,
					RemotePath: remotePath,
					Error:      err.Error(),
				}
			} else {
				results[idx] = *result
			}
		}(i, cfg)
	}
	wg.Wait()

	return results
}

// ============================================
// 内部辅助函数
// ============================================

// connect 建立 SSH 连接
func connect(cfg SSHConfig) (*ssh.Client, error) {
	// 默认值
	if cfg.Port == 0 {
		cfg.Port = 22
	}
	if cfg.User == "" {
		cfg.User = "root"
	}
	if cfg.Timeout == 0 {
		cfg.Timeout = 10 * time.Second
	}

	// 构建认证方法
	var authMethods []ssh.AuthMethod
	switch cfg.AuthType {
	case AuthKey:
		signer, err := ssh.ParsePrivateKey([]byte(cfg.KeyData))
		if err != nil {
			// 密钥解析失败，尝试用密码兜底
			if cfg.Password != "" {
				authMethods = []ssh.AuthMethod{ssh.Password(cfg.Password)}
			} else {
				return nil, fmt.Errorf("SSH密钥解析失败: %v", err)
			}
		} else {
			authMethods = []ssh.AuthMethod{ssh.PublicKeys(signer)}
		}
	default:
		authMethods = []ssh.AuthMethod{ssh.Password(cfg.Password)}
	}

	sshConfig := &ssh.ClientConfig{
		User:            cfg.User,
		Auth:            authMethods,
		Timeout:         cfg.Timeout,
		HostKeyCallback: tofuHostKeyCallback(cfg.Host),
	}

	addr := net.JoinHostPort(cfg.Host, fmt.Sprintf("%d", cfg.Port))
	client, err := ssh.Dial("tcp", addr, sshConfig)
	if err != nil {
		return nil, fmt.Errorf("SSH连接 %s 失败: %w", addr, err)
	}

	return client, nil
}

// tofuHostKeyCallback Trust-On-First-Use 主机密钥验证。
// 首次连接记录 fingerprint，后续连接校验。
func tofuHostKeyCallback(host string) ssh.HostKeyCallback {
	var knownKey ssh.PublicKey
	var mu sync.Mutex

	return func(hostname string, remote net.Addr, key ssh.PublicKey) error {
		mu.Lock()
		defer mu.Unlock()

		fp := ssh.FingerprintSHA256(key)

		if knownKey == nil {
			// 首次连接 — 信任并记录
			knownKey = key
			log.Printf("[RemoteExec] TOFU: trusting new host key for %s (%s): %s", host, hostname, fp)
			return nil
		}

		// 后续连接 — 校验
		if ssh.FingerprintSHA256(knownKey) == fp {
			return nil
		}

		return fmt.Errorf("SSH主机密钥变更！%s: expected %s, got %s",
			hostname, ssh.FingerprintSHA256(knownKey), fp)
	}
}

// FormatExecResults 格式化批量执行结果为可读字符串
func FormatExecResults(results []ExecResult) string {
	var buf bytes.Buffer
	successCount, failCount := 0, 0

	for _, r := range results {
		if r.Success {
			successCount++
		} else {
			failCount++
		}
	}

	buf.WriteString(fmt.Sprintf("==== 批量执行结果: %d 成功 / %d 失败 ====\n\n", successCount, failCount))

	for _, r := range results {
		status := "✓ 成功"
		if !r.Success {
			status = "✗ 失败"
		}
		buf.WriteString(fmt.Sprintf("── %s:%d [%s] (%dms)\n", r.Host, r.Port, status, r.Duration))
		if r.Error != "" {
			buf.WriteString(fmt.Sprintf("   错误: %s\n", r.Error))
		}
		if r.Output != "" && len(r.Output) < 2000 {
			buf.WriteString(fmt.Sprintf("   输出:\n%s\n", indent(r.Output, "     ")))
		} else if r.Output != "" {
			buf.WriteString(fmt.Sprintf("   输出(截断):\n%s\n...(共 %d 字符)\n", 
				r.Output[:1000], len(r.Output)))
		}
		buf.WriteString("\n")
	}

	return buf.String()
}

func indent(s, prefix string) string {
	lines := strings.Split(s, "\n")
	for i, line := range lines {
		lines[i] = prefix + line
	}
	return strings.Join(lines, "\n")
}