membank/auto-check/pkg/sshclient/sshclient.go

package sshclient

import (
	"fmt"
	"io"
	"net"
	"os"
	"path"
	"strings"
	"time"

	"github.com/pkg/sftp"
	"golang.org/x/crypto/ssh"
)

// Config SSH 连接配置（不含连接状态）
type Config struct {
	User     string
	Password string
	KeyFile  string
	Port     int
	Timeout  time.Duration
	SudoPwd  string // 非 root 用户执行特权命令(sudo)时的密码
}

// NewConfig 创建配置
func NewConfig(user, password, keyFile string, port int, timeout time.Duration) *Config {
	return &Config{
		User:     user,
		Password: password,
		KeyFile:  keyFile,
		Port:     port,
		Timeout:  timeout,
	}
}

// Connect 建立 SSH 连接，由调用方负责 Close
func Connect(ip string, cfg *Config) (*ssh.Client, error) {
	sshCfg, err := cfg.buildSSHConfig()
	if err != nil {
		return nil, err
	}

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

	// KeepAlive：长时压测（如 fnstable 门禁的 memtester 可达 30-120 分钟）期间
	// 连接上无数据流动，定期发 keepalive 请求防止 NAT/防火墙掐断空闲 TCP 连接。
	// 连接关闭后 SendRequest 返回错误，goroutine 自行退出。
	go func() {
		ticker := time.NewTicker(30 * time.Second)
		defer ticker.Stop()
		for range ticker.C {
			if _, _, err := client.SendRequest("keepalive@golang.org", true, nil); err != nil {
				return
			}
		}
	}()

	return client, nil
}

// RunCommand 在已有 SSH 连接上执行命令，返回输出（实现 stress.Executor 接口）
func RunCommand(client *ssh.Client, command string) (string, error) {
	session, err := client.NewSession()
	if err != nil {
		return "", fmt.Errorf("创建会话失败: %w", err)
	}
	defer session.Close()

	out, err := session.CombinedOutput(command)
	if err != nil {
		return cleanOutput(out), fmt.Errorf("命令执行失败: %w", err)
	}
	return cleanOutput(out), nil
}

// isNoiseLine 判断某行是否为登录横幅/欢迎语等噪声（fnOS 的 admin 用户会把
// "Could not chdir to home directory" 之类写入 stdout，污染命令解析）。
func isNoiseLine(line string) bool {
	s := strings.TrimSpace(line)
	if s == "" {
		return false
	}
	prefixes := []string{
		"Could not chdir to",
		"Last login:",
		"Welcome to",
		"/etc/motd",
		"Permission denied",
		"bash: line 1:",
	}
	for _, p := range prefixes {
		if strings.HasPrefix(s, p) {
			return true
		}
	}
	// 某些横幅是 "Could: command not found" 这类由前面噪声行产生，
	// 形如 "xxx: command not found" 且前半段是噪声关键词的也过滤
	if strings.Contains(s, ": command not found") {
		return true
	}
	return false
}

// cleanOutput 过滤掉登录横幅等噪声行，仅保留真实命令输出
func cleanOutput(raw []byte) string {
	lines := strings.Split(string(raw), "\n")
	var kept []string
	for _, ln := range lines {
		if isNoiseLine(ln) {
			continue
		}
		kept = append(kept, ln)
	}
	return strings.Join(kept, "\n")
}

// EnsureHome 确保远端登录用户的 HOME 目录存在，消除 fnOS "Could not chdir to
// home directory" 登录横幅（该横幅会写入 stdout 污染命令输出）。
func EnsureHome(client *ssh.Client) {
	out, err := RunCommand(client, "echo $HOME")
	home := strings.TrimSpace(out)
	if err != nil || home == "" || home == "/" {
		return
	}
	RunCommand(client, fmt.Sprintf("mkdir -p %s", home))
}

// UploadFile 通过 SFTP 把本地文件上传到远端指定路径（复用已有 SSH 连接）
func UploadFile(client *ssh.Client, localPath, remotePath string) error {
	src, err := os.Open(localPath)
	if err != nil {
		return fmt.Errorf("打开本地文件失败: %w", err)
	}
	defer src.Close()

	info, err := src.Stat()
	if err != nil {
		return fmt.Errorf("读取文件信息失败: %w", err)
	}

	sftpCli, err := sftp.NewClient(client)
	if err != nil {
		return fmt.Errorf("创建 SFTP 会话失败: %w", err)
	}
	defer sftpCli.Close()

	// 确保远端目录存在（/tmp 一般已存在，兜底处理）
	if dir := path.Dir(remotePath); dir != "" && dir != "/" {
		_ = sftpCli.MkdirAll(dir)
	}

	dst, err := sftpCli.Create(remotePath)
	if err != nil {
		return fmt.Errorf("创建远端文件失败: %w", err)
	}
	defer dst.Close()

	fmt.Printf("  [上传] %s -> %s (%.1f MB)...\n", localPath, remotePath, float64(info.Size())/1024/1024)
	start := time.Now()
	if _, err := io.Copy(dst, src); err != nil {
		return fmt.Errorf("文件传输失败: %w", err)
	}
	fmt.Printf("  [上传] 完成，耗时 %.1fs\n", time.Since(start).Seconds())
	return nil
}

// buildSSHConfig 构建认证配置
func (c *Config) buildSSHConfig() (*ssh.ClientConfig, error) {
	var authMethods []ssh.AuthMethod

	if c.Password != "" {
		authMethods = append(authMethods, ssh.Password(c.Password))
	}

	if c.KeyFile != "" {
		key, err := os.ReadFile(c.KeyFile)
		if err != nil {
			return nil, fmt.Errorf("读取密钥文件失败: %w", err)
		}
		signer, err := ssh.ParsePrivateKey(key)
		if err != nil {
			return nil, fmt.Errorf("解析私钥失败: %w", err)
		}
		authMethods = append(authMethods, ssh.PublicKeys(signer))
	}

	if len(authMethods) == 0 {
		return nil, fmt.Errorf("未提供认证方式（密码或密钥）")
	}

	return &ssh.ClientConfig{
		User:            c.User,
		Auth:            authMethods,
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		Timeout:         c.Timeout,
	}, nil
}