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

package config

import (
	"fmt"
	"os"
	"time"

	"gopkg.in/yaml.v3"
)

// ============================
// Config — 分类配置
// ============================

// Config 全局运行配置（仅从配置文件加载）
type Config struct {
	Scan     ScanConfig     `yaml:"scan"`
	SSH      SSHConfig      `yaml:"ssh"`
	Stress   StressConfig   `yaml:"stress"`
	Report   ReportConfig   `yaml:"report"`
	Workflow WorkflowConfig `yaml:"workflow"`
}

// ============================
// 分类配置项
// ============================

// ScanConfig 目标扫描配置（支持 CIDR 和 IP 范围格式）
type ScanConfig struct {
	CIDR        string        `yaml:"cidr"`
	Timeout     time.Duration `yaml:"timeout"`
	Concurrency int           `yaml:"concurrency"` // 并发探测数，默认 30
}

// SSHConfig SSH 认证
type SSHConfig struct {
	User     string `yaml:"user"`
	Password string `yaml:"password"`
	KeyFile  string `yaml:"key"`
	Port     int    `yaml:"port"`
	SudoPwd  string `yaml:"sudo_password"` // 非 root 用户执行特权命令时的 sudo 密码
}

// StressConfig 压力测试
type StressConfig struct {
	Types       string        `yaml:"types"`
	Duration    time.Duration `yaml:"duration"`
	Threads     int           `yaml:"threads"`
	MemSizeMB   int           `yaml:"mem_size_mb"`  // 内存压测占用(MB)，0=自动(可用60%)
	DiskSizeMB  int           `yaml:"disk_size_mb"` // 磁盘压测文件大小(MB)，0=默认1024
	DiskDir     string        `yaml:"disk_dir"`     // 磁盘压测目录，空=自动临时目录
	TempInterval time.Duration `yaml:"temp_interval"` // 温度采样间隔，0=默认10s
	TempLimit   float64       `yaml:"temp_limit"`   // 温度上限(°C)，0=不判定
	FnMemory    bool          `yaml:"fn_memory"`    // fnstable 门禁的 memtester 内存校验（30-120 分钟）
	// 压测统一在 Docker 容器中执行（工具与脚本预装在镜像里）
	DockerImage string `yaml:"docker_image"` // 压测镜像名称，如 stress-test:latest
	DockerTar   string `yaml:"docker_tar"`   // 本地导出的镜像文件（docker save），目标机无镜像时上传并 load
	DockerArgs  string `yaml:"docker_args"`  // Docker 运行参数，如 "--network=host --privileged"
}

// ReportConfig 报告
type ReportConfig struct {
	Path  string      `yaml:"path"`  // 报告保存路径
	Print PrintConfig `yaml:"print"` // 打印配置
	HTML  HTMLConfig  `yaml:"html"`  // HTML 报告配置
}

// HTMLConfig HTML 报告配置
type HTMLConfig struct {
	BaseURL  string `yaml:"base_url"`  // HTML 报告访问基地址（如 http://192.168.1.10:8080），空则只保存本地
	Path     string `yaml:"path"`      // HTML 报告保存目录，空 = 与 txt 报告同目录
	Filename string `yaml:"filename"`  // 文件名模板，支持 {mac} {ip} {ts}，默认 report-{mac}-{ip}-{ts}.html
	// QRURL 热敏打印中二维码指向的 URL 模板，支持 {filename} {mac} {ip} {ts} 占位
	// 留空时用 BaseURL + 实际文件名拼接
	QRURL string `yaml:"qr_url"`
}

// PrintConfig 打印配置
type PrintConfig struct {
	Enabled bool   `yaml:"enabled"` // 是否打印
	Backend string `yaml:"backend"` // system / escpos / usb / file
	Device  string `yaml:"device"`  // 目标：打印机名 / host:port / 设备路径 / 文件路径
	Width   int    `yaml:"width"`   // 行宽，80mm=42
}

// WorkflowConfig 工作流
type WorkflowConfig struct {
	Interval time.Duration `yaml:"interval"` // 轮询间隔
}

// ============================
// 加载
// ============================

// Default 内置默认配置
func Default() *Config {
	return &Config{
		Scan: ScanConfig{
			CIDR:        "192.168.1.0/24",
			Timeout:     3 * time.Second,
			Concurrency: 10,
		},
		SSH: SSHConfig{
			User: "root",
			Port: 22,
		},
		Stress: StressConfig{
			Types:        "cpu,memory",
			Duration:     30 * time.Second,
			Threads:      4,
			DiskSizeMB:   1024,
			TempInterval: 10 * time.Second,
			TempLimit:    90,
			FnMemory:     true,
			DockerImage:  "stress-test:latest",
			DockerTar:    "docker/stress-test/stress-test.tar",
			DockerArgs:   "--network=host --privileged",
		},
		Workflow: WorkflowConfig{
			Interval: 10 * time.Second,
		},
	}
}

// Load 从文件加载配置（覆盖默认值）
// path 为空时尝试加载 ./auto-check.yaml
// 文件不存在时返回默认配置
func Load(path string) (*Config, error) {
	cfg := Default()

	if path == "" {
		path = "auto-check.yaml"
	}

	data, err := os.ReadFile(path)
	if err != nil {
		if os.IsNotExist(err) {
			return cfg, nil // 无配置文件，使用默认值
		}
		return nil, fmt.Errorf("读取配置文件失败: %w", err)
	}

	if err := yaml.Unmarshal(data, cfg); err != nil {
		return nil, fmt.Errorf("解析配置文件 %s 失败: %w", path, err)
	}

	return cfg, nil
}

// LoadFromString 从字符串加载（用于测试）
func LoadFromString(s string) (*Config, error) {
	cfg := Default()
	if err := yaml.Unmarshal([]byte(s), cfg); err != nil {
		return nil, fmt.Errorf("解析配置失败: %w", err)
	}
	return cfg, nil
}