320 lines
11 KiB
Go
320 lines
11 KiB
Go
package stress
|
||
|
||
import (
|
||
"fmt"
|
||
"strings"
|
||
|
||
"auto-check/pkg/sshclient"
|
||
|
||
"golang.org/x/crypto/ssh"
|
||
)
|
||
|
||
// ============================
|
||
// 国内镜像源
|
||
// ============================
|
||
|
||
// setupMirror 在安装前把对应包管理器的源切换为国内镜像（清华 TUNA）。
|
||
// useMirror=false 时直接跳过（先用系统自带源试装）。失败不致命，仅提示并继续用原源。
|
||
func setupMirror(client *ssh.Client, pm string, sudoPwd string, useMirror bool) {
|
||
if !useMirror {
|
||
return
|
||
}
|
||
const label = "清华 TUNA 镜像"
|
||
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", label)
|
||
}
|
||
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", label)
|
||
default:
|
||
// dnf/yum 等无通用一键切源命令，需手动配置 repo 文件，这里仅提示
|
||
fmt.Printf(" [镜像] %s 无内置切源逻辑，建议使用系统默认源或手动配置%s repo\n", pm, label)
|
||
}
|
||
}
|
||
|
||
// ============================
|
||
// Docker 容器化压测支持
|
||
// ============================
|
||
|
||
// DockerInfo 远程 Docker 环境信息
|
||
type DockerInfo struct {
|
||
Available bool // Docker 是否可用（已安装且有权限）
|
||
Image string // 镜像名（检测时传入，确认存在时回填）
|
||
HasImage bool // 目标镜像是否已存在本地（避免每次 pull）
|
||
}
|
||
|
||
// DetectDocker 检测远程主机的 Docker 环境。
|
||
// 包括：docker 命令是否存在、是否需要 sudo、镜像是否已存在。
|
||
// sudoPwd 为非 root 用户执行时的密码（用于 sudo docker）。
|
||
func DetectDocker(client *ssh.Client, image string, sudoPwd string) DockerInfo {
|
||
info := DockerInfo{Image: image}
|
||
|
||
// 1. 检测 docker 命令是否存在（支持 root 和 sudo）
|
||
if out, err := runPrivileged(client, "command -v docker", sudoPwd); err != nil || strings.TrimSpace(out) == "" {
|
||
return info
|
||
}
|
||
info.Available = true
|
||
|
||
// 2. 检测 docker 服务是否运行（docker info 成功即可）
|
||
if _, err := runPrivileged(client, "docker info >/dev/null 2>&1", sudoPwd); err != nil {
|
||
info.Available = false
|
||
return info
|
||
}
|
||
|
||
// 3. 检测目标镜像是否已存在本地（避免每次拉取）
|
||
if out, err := runPrivileged(client, fmt.Sprintf("docker image inspect %s >/dev/null 2>&1 && echo EXIST", image), sudoPwd); err == nil && strings.Contains(out, "EXIST") {
|
||
info.HasImage = true
|
||
}
|
||
|
||
return info
|
||
}
|
||
|
||
// ============================
|
||
// 远程工具检测
|
||
// ============================
|
||
|
||
// requiredTools 压测核心依赖工具（缺则按需在线安装）。lm-sensors 依赖较多，
|
||
// 仅单文件二进制 stress-ng 与 fio 为必需；sensors 缺失仅影响温度曲线，不阻断压测。
|
||
// 这是工具检测的唯一来源，DetectTools/EnsureTools 都从这里取列表，避免两处不一致。
|
||
var requiredTools = []string{"stress-ng", "fio"}
|
||
|
||
// ToolInfo 远程工具信息
|
||
type ToolInfo struct {
|
||
Name string
|
||
Path string
|
||
Version string
|
||
}
|
||
|
||
// ToolSet 远程可用工具集合
|
||
type ToolSet struct {
|
||
Tools map[string]ToolInfo // name → info
|
||
}
|
||
|
||
// Has 工具是否可用
|
||
func (ts ToolSet) Has(name string) bool {
|
||
_, ok := ts.Tools[name]
|
||
return ok
|
||
}
|
||
|
||
// Get 获取工具信息
|
||
func (ts ToolSet) Get(name string) (ToolInfo, bool) {
|
||
t, ok := ts.Tools[name]
|
||
return t, ok
|
||
}
|
||
|
||
// DetectTools 检测远程所有相关工具是否可用
|
||
func DetectTools(client *ssh.Client) ToolSet {
|
||
ts := ToolSet{Tools: make(map[string]ToolInfo)}
|
||
for _, name := range requiredTools {
|
||
if info, ok := detectOne(client, name); ok {
|
||
ts.Tools[name] = info
|
||
}
|
||
}
|
||
return ts
|
||
}
|
||
|
||
// EnsureTools 检测工具；缺失且 autoInstall=true 时按发行版自动安装。
|
||
// sudoPwd 为非 root 用户执行安装时提供（为空则尝试免密 sudo）。useMirror=true
|
||
// 时安装前把包管理器源切换为国内镜像（清华 TUNA）。返回安装后的工具集合，
|
||
// 以及最终仍缺失的工具名列表。
|
||
func EnsureTools(client *ssh.Client, names []string, autoInstall bool, sudoPwd string, useMirror bool) (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, useMirror)
|
||
// 仅 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) {
|
||
path := ""
|
||
if out, err := sshclient.RunCommand(client, "which "+name+" 2>/dev/null"); err == nil {
|
||
path = strings.TrimSpace(out)
|
||
}
|
||
if path == "" {
|
||
return ToolInfo{}, false
|
||
}
|
||
|
||
// 取版本：lm-sensors 仅支持 -v，其余工具均支持 --version
|
||
verFlag := "--version"
|
||
if name == "lm-sensors" {
|
||
verFlag = "-v"
|
||
}
|
||
v, _ := sshclient.RunCommand(client, path+" "+verFlag+" 2>&1 | head -1")
|
||
version := strings.TrimSpace(v)
|
||
|
||
return ToolInfo{Name: name, Path: path, Version: version}, true
|
||
}
|