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

package stress

import (
	"fmt"
	"strings"
	"time"

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

// ============================
// 远程命令执行接口
// ============================

// Executor 远程命令执行器
type Executor interface {
	RunCommand(client *ssh.Client, command string) (string, error)
}

// exec 执行远程命令（带超时）
func exec(client *ssh.Client, executor Executor, cmd string) (string, error) {
	return executor.RunCommand(client, cmd)
}

// execf 格式化执行
func execf(client *ssh.Client, executor Executor, format string, args ...interface{}) (string, error) {
	return exec(client, executor, fmt.Sprintf(format, args...))
}

// ============================
// 工具检测
// ============================

// ToolInfo 远程工具信息
type ToolInfo struct {
	Available bool
	Path      string
	Version   string
}

// DetectTool 检测远程工具是否可用
func DetectTool(client *ssh.Client, executor Executor, name string) ToolInfo {
	// 检查路径
	out, err := execf(client, executor, "which %s 2>/dev/null", name)
	path := strings.TrimSpace(out)
	if err != nil || path == "" {
		return ToolInfo{Available: false}
	}

	// 获取版本
	var version string
	switch name {
	case "stress-ng":
		out, _ = execf(client, executor, "%s --version 2>&1 | head -1", path)
		version = strings.TrimSpace(out)
	case "stressapptest":
		out, _ = execf(client, executor, "%s --help 2>&1 | head -1", path)
		version = strings.TrimSpace(out)
	case "lm-sensors":
		out, _ = execf(client, executor, "%s -v 2>&1 | head -1", path)
		version = strings.TrimSpace(out)
	}

	return ToolInfo{Available: true, Path: path, Version: version}
}

// EnsureTool 检测工具，不可用则返回错误提示
func EnsureTool(client *ssh.Client, executor Executor, name string, installHint string) (ToolInfo, error) {
	info := DetectTool(client, executor, name)
	if !info.Available {
		return info, fmt.Errorf("%s 未安装。安装方式: %s", name, installHint)
	}
	return info, nil
}

// ============================
// 环境探测
// ============================

// SystemInfo 系统基础信息
type SystemInfo struct {
	Hostname  string
	CPUModel  string
	CPUCores  string
	MemTotal  string
	KernelVer string
	DiskInfo  string
}

// ProbeSystem 探测远程系统基础信息
func ProbeSystem(client *ssh.Client, executor Executor) SystemInfo {
	info := SystemInfo{}

	out, _ := exec(client, executor, "hostname 2>/dev/null")
	info.Hostname = strings.TrimSpace(out)

	out, _ = exec(client, executor, "lscpu 2>/dev/null | grep 'Model name' | sed 's/Model name:\\s*//'")
	info.CPUModel = strings.TrimSpace(out)

	out, _ = exec(client, executor, "nproc 2>/dev/null")
	info.CPUCores = strings.TrimSpace(out)

	out, _ = exec(client, executor, "free -h 2>/dev/null | awk '/Mem:/{print $2}'")
	info.MemTotal = strings.TrimSpace(out)

	out, _ = exec(client, executor, "uname -r 2>/dev/null")
	info.KernelVer = strings.TrimSpace(out)

	out, _ = exec(client, executor, "lsblk -d -o NAME,SIZE,TYPE 2>/dev/null | head -5")
	info.DiskInfo = strings.TrimSpace(out)

	return info
}

// ============================
// 温度监控
// ============================

// TempMonitor 温度监控器
type TempMonitor struct {
	client   *ssh.Client
	executor Executor
	interval time.Duration
	stopCh   chan struct{}
	logs     []string
}

// NewTempMonitor 创建温度监控器
func NewTempMonitor(client *ssh.Client, executor Executor, interval time.Duration) *TempMonitor {
	if interval <= 0 {
		interval = 10 * time.Second
	}
	return &TempMonitor{
		client:   client,
		executor: executor,
		interval: interval,
		stopCh:   make(chan struct{}),
	}
}

// Start 后台启动温度监控
func (m *TempMonitor) Start() {
	go m.loop()
}

// Stop 停止监控并返回所有采样
func (m *TempMonitor) Stop() []string {
	close(m.stopCh)
	return m.logs
}

func (m *TempMonitor) loop() {
	ticker := time.NewTicker(m.interval)
	defer ticker.Stop()

	// 初始温度
	m.sample()

	for {
		select {
		case <-m.stopCh:
			return
		case <-ticker.C:
			m.sample()
		}
	}
}

func (m *TempMonitor) sample() {
	// 尝试 sensors
	out, err := execf(m.client, m.executor, "sensors 2>/dev/null | grep -i 'temp\\|core\\|cpu' | head -10")
	if err == nil && strings.TrimSpace(out) != "" {
		ts := time.Now().Format("15:04:05")
		for _, line := range strings.Split(out, "\n") {
			line = strings.TrimSpace(line)
			if line != "" {
				m.logs = append(m.logs, fmt.Sprintf("[%s] %s", ts, line))
			}
		}
		return
	}

	// 回退: 读 sysfs
	out, _ = execf(m.client, m.executor, `cat /sys/class/thermal/thermal_zone*/temp 2>/dev/null | while read t; do echo "$((t/1000))°C"; done`)
	if strings.TrimSpace(out) != "" {
		ts := time.Now().Format("15:04:05")
		m.logs = append(m.logs, fmt.Sprintf("[%s] %s", ts, strings.TrimSpace(out)))
	}
}

// ============================
// dmesg 监控
// ============================

// DmesgMonitor 内核日志监控
type DmesgMonitor struct {
	client   *ssh.Client
	executor Executor
	stopCh   chan struct{}
	baseline string // 启动时的 dmesg 行数
	logs     []string
}

// NewDmesgMonitor 创建 dmesg 监控器
func NewDmesgMonitor(client *ssh.Client, executor Executor) *DmesgMonitor {
	return &DmesgMonitor{
		client:   client,
		executor: executor,
		stopCh:   make(chan struct{}),
	}
}

// Start 记录基线，后台监控
func (m *DmesgMonitor) Start() {
	out, _ := exec(m.client, m.executor, "dmesg 2>/dev/null | wc -l")
	m.baseline = strings.TrimSpace(out)
	go m.loop()
}

// Stop 停止监控，返回新增的硬件报错
func (m *DmesgMonitor) Stop() []string {
	close(m.stopCh)

	// 获取新增的 dmesg 日志中的硬件错误
	out, _ := execf(m.client, m.executor,
		`dmesg --level=err,crit,alert,emerg 2>/dev/null | tail -30`)
	if strings.TrimSpace(out) != "" {
		m.logs = append(m.logs, strings.Split(out, "\n")...)
	}

	// 过滤硬件相关关键词
	var hwErrors []string
	keywords := []string{"error", "fail", "fault", "warn", "critical", "oom", "panic", "hardware", "thermal", "throttl"}
	for _, line := range m.logs {
		lower := strings.ToLower(line)
		for _, kw := range keywords {
			if strings.Contains(lower, kw) {
				hwErrors = append(hwErrors, strings.TrimSpace(line))
				break
			}
		}
	}
	return hwErrors
}

func (m *DmesgMonitor) loop() {
	ticker := time.NewTicker(5 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-m.stopCh:
			return
		case <-ticker.C:
			// 周期性检查是否有新的硬件相关日志（轻量级）
			out, _ := execf(m.client, m.executor,
				`dmesg --level=err,crit 2>/dev/null | tail -3`)
			if strings.TrimSpace(out) != "" {
				ts := time.Now().Format("15:04:05")
				for _, line := range strings.Split(out, "\n") {
					line = strings.TrimSpace(line)
					if line != "" {
						m.logs = append(m.logs, fmt.Sprintf("[%s] %s", ts, line))
					}
				}
			}
		}
	}
}