membank/auto-check/pkg/stress/runner.go

package stress

import (
	"fmt"
	"strconv"
	"strings"
	"time"

	"auto-check/pkg/config"
	"auto-check/pkg/sshclient"

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

// Results 测试结果表（IP → 报告），包级公开
var Results = make(map[string]*Report)

// ============================
// 执行器
// ============================

// Runner 压力测试执行器
type Runner struct {
	cfg     Config
	types   []TestType
	client  *ssh.Client
	tools   ToolSet
	metrics []Metric
}

// NewRunner 创建执行器（检测工具 → 构建指标 → 生成脚本 → 执行）
func NewRunner(cfg Config, types []TestType, client *ssh.Client) *Runner {
	tools := DetectTools(client)
	metrics := BuildMetrics(types, tools)
	return &Runner{cfg: cfg, types: types, client: client, tools: tools, metrics: metrics}
}

// Run 执行压力测试（生成脚本 → 远程执行 → 解析结果）
func (r *Runner) Run(ip string) *Report {
	report := &Report{IP: ip, StartTime: time.Now()}
	fmt.Printf("\n════════ [%s] 压力测试开始 ════════\n", ip)

	// 显示可用工具
	fmt.Printf("  工具: ")
	for name, info := range r.tools.Tools {
		fmt.Printf("%s(%s) ", name, info.Version)
	}
	fmt.Println()

	// 显示将执行的指标
	fmt.Printf("  指标: ")
	for _, m := range r.metrics {
		if m.Enabled {
			fmt.Printf("%s ", m.Name)
		}
	}
	fmt.Println()

	// 生成脚本
	script := BuildScript(r.metrics, r.cfg)

	// 远程执行脚本（通过 stdin 传入）
	fmt.Printf("  [执行] 生成脚本 %d 字节，开始远程执行...\n", len(script))
	scriptCmd := fmt.Sprintf("cat <<'AUTOCHECKSCRIPT' > /tmp/auto-check.sh\n%s\nAUTOCHECKSCRIPT\nchmod +x /tmp/auto-check.sh && bash /tmp/auto-check.sh 2>&1", script)
	output, err := sshclient.RunCommand(r.client, scriptCmd)

	if err != nil {
		fmt.Printf("  [执行] 脚本执行出错: %v\n", err)
	}

	// 解析结果
	if output != "" {
		r.parseResults(report, output)
	}

	report.EndTime = time.Now()
	report.Duration = report.EndTime.Sub(report.StartTime)
	fmt.Printf("\n════════ [%s] 压力测试完成 ════════\n", ip)
	fmt.Println(report.ToText())
	return report
}

// parseResults 解析脚本的结构化输出
func (r *Runner) parseResults(report *Report, output string) {
	lines := strings.Split(output, "\n")
	var currentTest string
	var currentStatus string
	var currentDuration string
	var currentOutput []string

	flush := func() {
		if currentTest == "" {
			return
		}
		status := StatusPass
		switch currentStatus {
		case "fail":
			status = StatusFail
		case "skip":
			status = StatusSkip
		case "error":
			status = StatusError
		}
		duration, _ := time.ParseDuration(currentDuration)
		report.AddResult(Result{
			Type:     TestType(currentTest),
			Status:   status,
			Output:   strings.Join(currentOutput, "\n"),
			Duration: duration,
		})
		currentTest = ""
		currentStatus = ""
		currentDuration = ""
		currentOutput = nil
	}

	for _, line := range lines {
		line = strings.TrimSpace(line)

		if strings.HasPrefix(line, "===TEST:") {
			flush()
			currentTest = strings.TrimSuffix(strings.TrimPrefix(line, "===TEST:"), "===")
			continue
		}

		if strings.HasPrefix(line, "===MONITOR:") {
			flush()
			currentTest = strings.TrimSuffix(strings.TrimPrefix(line, "===MONITOR:"), "===")
			continue
		}

		if line == "===END===" {
			flush()
			continue
		}

		// 解析 key:value
		if strings.HasPrefix(line, "status:") {
			currentStatus = strings.TrimPrefix(line, "status:")
			continue
		}
		if strings.HasPrefix(line, "duration:") {
			durStr := strings.TrimPrefix(line, "duration:")
			ms, _ := strconv.ParseInt(strings.TrimSuffix(durStr, "ms"), 10, 64)
			currentDuration = fmt.Sprintf("%dms", ms)
			continue
		}
		if strings.HasPrefix(line, "output:") {
			currentOutput = append(currentOutput, strings.TrimPrefix(line, "output:"))
			continue
		}
		if strings.HasPrefix(line, "samples:") || strings.HasPrefix(line, "max:") {
			currentOutput = append(currentOutput, line)
			continue
		}

		// 普通输出行
		if currentTest != "" && line != "" {
			currentOutput = append(currentOutput, line)
		}
	}
	flush()

	// 从系统信息输出解析主机信息
	for _, line := range lines {
		if strings.HasPrefix(line, "hostname:") {
			// 可扩展：存入 report 的系统信息字段
		}
	}
}

// ============================
// 业务入口（workflow 调用）
// ============================

// TestDevice 对单台设备执行 SSH 登录 + 压力测试（业务入口）
func TestDevice(ip string, cfg config.Config) *Report {
	sshCfg := sshclient.NewConfig(cfg.SSH.User, cfg.SSH.Password, cfg.SSH.KeyFile, cfg.SSH.Port, cfg.Scan.Timeout)

	conn, err := sshclient.Connect(ip, sshCfg)
	if err != nil {
		fmt.Printf("  [测试] %s SSH 连接失败: %v\n", ip, err)
		return NewSSHFailReport(ip, err)
	}
	defer conn.Close()

	stressCfg := Config{
		Duration:   cfg.Stress.Duration,
		Threads:    cfg.Stress.Threads,
		DiskSizeMB: 1024,
		TempLogInt: 10 * time.Second,
	}
	return NewRunner(stressCfg, ParseTypes(cfg.Stress.Types), conn).Run(ip)
}

// NewSSHFailReport 构造 SSH 连接失败报告
func NewSSHFailReport(ip string, err error) *Report {
	return &Report{
		IP:        ip,
		StartTime: time.Now(),
		Results:   []Result{{Type: "ssh", Status: StatusFail, Error: err.Error()}},
		Failed:    1,
	}
}

// ============================
// 工具函数
// ============================

// IsPassed 报告是否通过
func IsPassed(rpt *Report) bool {
	return rpt != nil && rpt.Failed == 0 && rpt.Errors == 0
}

// Status 报告状态字符串
func Status(rpt *Report) string {
	if rpt == nil {
		return StatusUnknown
	}
	if IsPassed(rpt) {
		return StatusPass
	}
	return StatusFail
}

// ParseTypes 逗号分隔字符串 → TestType 列表
func ParseTypes(s string) []TestType {
	var types []TestType
	for _, t := range strings.Split(s, ",") {
		if t = strings.TrimSpace(t); t != "" {
			types = append(types, TestType(t))
		}
	}
	return types
}