0813

2026-08-13 22:49:50 +08:00
parent 5e296bfe19
commit 2aaaa16ae2
24 changed files with 1415 additions and 327 deletions
--- a/auto-check/pkg/stress/tools.go
+++ b/auto-check/pkg/stress/tools.go
@@ -1,6 +1,7 @@
 package stress

 import (
+	"fmt"
 	"strings"

 	"auto-check/pkg/sshclient"
@@ -8,10 +9,82 @@ import (
 	"golang.org/x/crypto/ssh"
 )

+// ============================
+// 国内镜像源
+// ============================
+
+// useMirror 控制是否把远端包管理器切到国内镜像。默认关闭——先用系统自带源
+// 试装（飞牛/群晖等定制系统常常已预置可用源），失败或超时再开启清华镜像。
+var useMirror = false
+
+// SetMirror 设置是否使用国内镜像源（由配置注入）
+func SetMirror(on bool) {
+	useMirror = on
+}
+
+// mirrorName 根据发行版返回镜像别名（用于日志）
+const mirrorLabel = "清华 TUNA 镜像"
+
+// setupMirror 在安装前把对应包管理器的源切换为国内镜像（仅 apt 系需要改文件，
+// 其它包管理器通过参数/变量临时指向镜像）。失败不致命，仅打印警告并继续用原源。
+func setupMirror(client *ssh.Client, pm string, sudoPwd string) {
+	if !useMirror {
+		return
+	}
+	switch pm {
+	case "apt-get":
+		// 备份并替换为清华镜像（Debian/Ubuntu 通用，按 lsb_release 选版本）
+		script := `set -e
+OS_ID=$(. /etc/os-release 2>/dev/null; echo "$ID")
+OS_VER=$(. /etc/os-release 2>/dev/null; echo "$VERSION_CODENAME")
+[ -z "$OS_VER" ] && OS_VER=$(. /etc/os-release 2>/dev/null; echo "$VERSION_ID")
+MIRROR="https://mirrors.tuna.tsinghua.edu.cn"
+if [ "$OS_ID" = "ubuntu" ]; then
+  LIST="$MIRROR/ubuntu"
+elif [ "$OS_ID" = "debian" ]; then
+  LIST="$MIRROR/debian"
+elif [ "$OS_ID" = "raspbian" ]; then
+  LIST="$MIRROR/raspbian"
+else
+  LIST="$MIRROR/debian"
+fi
+BAK=/etc/apt/sources.list.d/auto-check-mirror.bak
+if [ ! -f "$BAK" ]; then
+  cp /etc/apt/sources.list "$BAK" 2>/dev/null || true
+  cp -r /etc/apt/sources.list.d "$BAK.d" 2>/dev/null || true
+fi
+: > /etc/apt/sources.list
+echo "deb $LIST $OS_VER main contrib non-free" > /etc/apt/sources.list
+echo "deb $LIST $OS_VER-updates main contrib non-free" >> /etc/apt/sources.list
+echo "deb $LIST $OS_VER-backports main contrib non-free" >> /etc/apt/sources.list
+echo "deb $LIST-security $OS_VER-security main contrib non-free" >> /etc/apt/sources.list
+echo OK`
+		out, err := runPrivileged(client, "sh -c "+shellQuote(script), sudoPwd)
+		if err != nil || !strings.Contains(out, "OK") {
+			fmt.Printf("  [镜像] apt 切换国内源失败，将使用官方源: %v\n", err)
+		} else {
+			fmt.Printf("  [镜像] apt 已切换至%s\n", mirrorLabel)
+		}
+	case "dnf", "yum":
+		// 清华镜像站对 dnf/yum 提供 repo 文件；这里用变量临时指向（持久化需改 repo）
+		cmd := `echo "腾讯/清华镜像需手动配置 repo；尝试临时使用 mirrors.tuna.tsinghua.edu.cn 变量"`
+		runPrivileged(client, cmd, sudoPwd)
+		// 通过 baseurl 变量临时覆盖（仅对部分仓库有效，作为加速兜底）
+		fmt.Printf("  [镜像] %s 使用系统默认源（建议手动配置%s repo）\n", pm, mirrorLabel)
+	case "apk":
+		runPrivileged(client, "sed -i 's#https\\?://[^/]*alpinelinux.org#https://mirrors.tuna.tsinghua.edu.cn/alpine#g' /etc/apk/repositories", sudoPwd)
+		fmt.Printf("  [镜像] apk 已切换至%s\n", mirrorLabel)
+	}
+}
+
 // ============================
 // 远程工具检测
 // ============================

+// 压测核心依赖工具（缺则按需在线安装）。lm-sensors 依赖较多，仅单文件
+// 二进制 stress-ng 与 fio 为必需；sensors 缺失仅影响温度曲线，不阻断压测。
+var requiredTools = []string{"stress-ng", "fio"}
+
 // ToolInfo 远程工具信息
 type ToolInfo struct {
 	Name    string
@@ -47,31 +120,185 @@ func DetectTools(client *ssh.Client) ToolSet {
 	return ts
 }

-// detectOne 检测单个工具
+// EnsureTools 检测工具；缺失且 autoInstall=true 时按发行版自动安装。
+// sudoPwd 为非 root 用户执行安装时提供（为空则尝试免密 sudo）。
+// 返回安装后的工具集合，以及最终仍缺失的工具名列表。
+func EnsureTools(client *ssh.Client, names []string, autoInstall bool, sudoPwd string) (ToolSet, []string) {
+	ts := DetectTools(client)
+	var missing []string
+	for _, n := range names {
+		if !ts.Has(n) {
+			missing = append(missing, n)
+		}
+	}
+	if len(missing) == 0 {
+		return ts, nil
+	}
+
+	var failed []string
+	if autoInstall {
+		pm := detectPkgManager(client)
+		if pm == "" {
+			return ts, missing // 无法识别包管理器，跳过安装
+		}
+		// 安装前先把源切到国内镜像（使用官方源在网络差的环境会超时）
+		setupMirror(client, pm, sudoPwd)
+		// 仅 apt 系需要先刷新元数据
+		if pm == "apt-get" {
+			if out, err := runPrivileged(client, "timeout 180 apt-get update 2>&1", sudoPwd); err != nil || strings.Contains(out, "Err") || strings.Contains(out, "Failed") {
+				fmt.Printf("  [安装] apt-get update 异常(可能无可用源/无权限): %v\n%s\n", err, truncate(out, 400))
+			}
+		}
+		for _, n := range missing {
+			pkg := pkgName(pm, n)
+			if pkg == "" {
+				failed = append(failed, n)
+				continue
+			}
+			installCmd := pkgInstallCmd(pm, pkg)
+			// runPrivileged 统一处理 sudo 与引号，这里直接给裸安装命令（带 timeout 兜底）
+			out, err := runPrivileged(client, fmt.Sprintf("timeout 300 %s 2>&1", installCmd), sudoPwd)
+			if err != nil || !strings.Contains(out, "Setting up "+pkg) && !strings.Contains(out, "已安装") && !strings.Contains(out, "already newest") {
+				fmt.Printf("  [安装] %s 安装失败: %v\n%s\n", pkg, err, truncate(out, 500))
+			}
+			// 安装后重新检测该工具是否到位
+			if info, ok := detectOne(client, n); ok {
+				ts.Tools[n] = info
+			} else {
+				failed = append(failed, n)
+			}
+		}
+		if len(failed) == 0 {
+			return ts, nil
+		}
+		return ts, failed
+	}
+	return ts, missing
+}
+
+// runPrivileged 以特权执行命令：root 直接执行；非 root 用 sudo，
+// 若提供了 sudo 密码则通过 stdin 喂入（无需 NOPASSWD 配置）。
+func runPrivileged(client *ssh.Client, cmd, sudoPwd string) (string, error) {
+	// 已是 root 则直接执行
+	if out, err := sshclient.RunCommand(client, "id -u"); err == nil && strings.TrimSpace(out) == "0" {
+		return sshclient.RunCommand(client, cmd)
+	}
+	// 非 root：尝试 sudo
+	if sudoPwd != "" {
+		// echo 密码 | sudo -S 执行；整条命令用双引号包裹，避免命令内部单引号嵌套冲突。
+		// 仅对命令内的双引号转义（安装命令均不含双引号，安全）。
+		quoted := strings.ReplaceAll(singleLine(cmd), `"`, `\"`)
+		wrapped := fmt.Sprintf("echo %s | sudo -S -- sh -c \"%s\"", shellQuote(sudoPwd), quoted)
+		return sshclient.RunCommand(client, wrapped)
+	}
+	// 无密码：先尝试免密 sudo (-n)，失败则回退普通 sudo（会提示需要密码）
+	quoted := strings.ReplaceAll(singleLine(cmd), `"`, `\"`)
+	if out, err := sshclient.RunCommand(client, "sudo -n -- sh -c \""+quoted+"\""); err == nil {
+		return out, nil
+	}
+	return sshclient.RunCommand(client, "sudo -- sh -c \""+quoted+"\"")
+}
+
+// shellQuote 转义单引号，避免密码破坏命令结构
+func shellQuote(s string) string {
+	return "'" + strings.ReplaceAll(s, "'", `'\''`) + "'"
+}
+
+// singleLine 去掉换行，保证能塞进 sh -c '...'
+func singleLine(s string) string {
+	return strings.ReplaceAll(strings.ReplaceAll(s, "\n", " "), "\r", "")
+}
+
+// truncate 截断过长的输出，便于日志展示
+func truncate(s string, n int) string {
+	// 去掉 ANSI/多余空行，只保留前 n 个字符
+	s = strings.TrimSpace(s)
+	if len([]rune(s)) > n {
+		return string([]rune(s)[:n]) + "...(truncated)"
+	}
+	return s
+}
+
+// detectPkgManager 识别远程发行版的包管理器
+func detectPkgManager(client *ssh.Client) string {
+	for _, pm := range []string{"apt-get", "dnf", "yum", "apk", "zypper"} {
+		if _, err := sshclient.RunCommand(client, fmt.Sprintf("command -v %s", pm)); err == nil {
+			return pm
+		}
+	}
+	return ""
+}
+
+// pkgName 工具名 → 对应发行版的包名
+func pkgName(pm, tool string) string {
+	m := map[string]string{
+		"stress-ng":     "stress-ng",
+		"fio":           "fio",
+		"lm-sensors":    "lm-sensors",
+		"stressapptest": "stressapptest",
+		"iperf3":        "iperf3",
+	}
+	if pm == "dnf" || pm == "yum" {
+		// RHEL 系 lm-sensors 包名为 lm_sensors
+		if tool == "lm-sensors" {
+			return "lm_sensors"
+		}
+	}
+	if pm == "zypper" {
+		if tool == "lm-sensors" {
+			return "sensors"
+		}
+	}
+	if v, ok := m[tool]; ok {
+		return v
+	}
+	return tool
+}
+
+// pkgInstallCmd 生成安装命令
+func pkgInstallCmd(pm, pkg string) string {
+	switch pm {
+	case "apt-get":
+		return fmt.Sprintf("apt-get install -y %s", pkg)
+	case "dnf":
+		return fmt.Sprintf("dnf install -y %s", pkg)
+	case "yum":
+		return fmt.Sprintf("yum install -y %s", pkg)
+	case "apk":
+		return fmt.Sprintf("apk add %s", pkg)
+	case "zypper":
+		return fmt.Sprintf("zypper install -y %s", pkg)
+	}
+	return ""
+}
+
+// detectOne 检测单个工具（查系统 PATH；缺失由 EnsureTools 在线安装补全）
 func detectOne(client *ssh.Client, name string) (ToolInfo, bool) {
-	out, err := sshclient.RunCommand(client, "which "+name+" 2>/dev/null")
-	path := strings.TrimSpace(out)
-	if err != nil || path == "" {
+	path := ""
+	if out, err := sshclient.RunCommand(client, "which "+name+" 2>/dev/null"); err == nil {
+		path = strings.TrimSpace(out)
+	}
+	if path == "" {
 		return ToolInfo{}, false
 	}

 	var version string
 	switch name {
 	case "stress-ng":
-		out, _ = sshclient.RunCommand(client, path+" --version 2>&1 | head -1")
-		version = strings.TrimSpace(out)
+		v, _ := sshclient.RunCommand(client, path+" --version 2>&1 | head -1")
+		version = strings.TrimSpace(v)
 	case "stressapptest":
-		out, _ = sshclient.RunCommand(client, path+" --help 2>&1 | head -1")
-		version = strings.TrimSpace(out)
+		v, _ := sshclient.RunCommand(client, path+" --help 2>&1 | head -1")
+		version = strings.TrimSpace(v)
 	case "iperf3":
-		out, _ = sshclient.RunCommand(client, path+" --version 2>&1 | head -1")
-		version = strings.TrimSpace(out)
+		v, _ := sshclient.RunCommand(client, path+" --version 2>&1 | head -1")
+		version = strings.TrimSpace(v)
 	case "fio":
-		out, _ = sshclient.RunCommand(client, path+" --version 2>&1 | head -1")
-		version = strings.TrimSpace(out)
+		v, _ := sshclient.RunCommand(client, path+" --version 2>&1 | head -1")
+		version = strings.TrimSpace(v)
 	case "lm-sensors":
-		out, _ = sshclient.RunCommand(client, path+" -v 2>&1 | head -1")
-		version = strings.TrimSpace(out)
+		v, _ := sshclient.RunCommand(client, path+" -v 2>&1 | head -1")
+		version = strings.TrimSpace(v)
 	}

 	return ToolInfo{Name: name, Path: path, Version: version}, true