142 lines
3.9 KiB
Go
142 lines
3.9 KiB
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=不判定
|
||
AutoInstall bool `yaml:"auto_install"` // 缺工具时尝试在线安装（apt/dnf/yum 等）
|
||
UseMirror bool `yaml:"use_mirror"` // 在线安装时切换国内镜像源（清华 TUNA），国内网络卡顿/超时再开启
|
||
}
|
||
|
||
// ReportConfig 报告
|
||
type ReportConfig struct {
|
||
Path string `yaml:"path"` // 报告保存路径
|
||
Print PrintConfig `yaml:"print"` // 打印配置
|
||
}
|
||
|
||
// 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,
|
||
AutoInstall: true,
|
||
UseMirror: false,
|
||
},
|
||
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
|
||
}
|