membank/auto-check/pkg/workflow/loop.go

package workflow

import (
	"fmt"
	"strings"
	"time"

	"auto-check/pkg/discovery"
	"auto-check/pkg/report"
	"auto-check/pkg/sshclient"
	"auto-check/pkg/stress"
)

// ============================
// 状态常量
// ============================

const (
	StatusUntested = "untested" // 未测试
	StatusPass     = "pass"     // 通过
	StatusFail     = "fail"     // 失败
)

// ============================
// 设备（MAC 为唯一标识）
// ============================

// Device 设备信息
type Device struct {
	MAC  string
	IP   string
	Host discovery.Host
}

// ============================
// 永久循环工作流
// ============================

// Loop 永久循环：每间隔执行一轮
// 每轮: 扫描设备 → 增量测试 → 打印状态变化
func Loop(ctx *Context) error {
	interval := ctx.Config.Workflow.Interval
	if interval <= 0 {
		interval = 10 * time.Second
	}

	fmt.Printf("══ 工作流启动（永久循环，间隔 %v，设备以 MAC 唯一标识）══\n\n", interval)

	// 第一轮立即执行
	round(ctx)

	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			fmt.Println("══ 工作流已停止 ══")
			return nil
		case <-ticker.C:
			round(ctx)
		}
	}
}

// round 执行一轮
func round(ctx *Context) {
	now := time.Now().Format("15:04:05")
	fmt.Printf("════ 轮次开始 [%s] ════\n", now)

	// ── 1. 扫描设备，更新设备列表 map（MAC 为 key）──
	scanDevices(ctx)

	// ── 2. 遍历设备: 结果为空或失败 → 开始测试，保存报告和结果 ──
	testPending(ctx)

	// ── 3. 遍历打印: 状态变化才重新打印 ──
	printStatus(ctx)
}

// scanDevices 扫描并更新设备列表
// 设备身份 = MAC；IP 变化视为同一设备（更新 IP），IP 被复用则移除旧设备
func scanDevices(ctx *Context) {
	scanner := discovery.NewScanner(
		ctx.Config.Scan.CIDR,
		ctx.Config.SSH.Port,
		ctx.Config.Scan.Timeout,
		ctx.Config.Scan.Concurrency,
	)
	hosts := scanner.Scan()

	// 本轮扫描到的 IP → MAC 映射
	ipToMAC := make(map[string]string, len(hosts))
	for _, h := range hosts {
		ipToMAC[h.IP.String()] = h.MAC
	}

	// 1. 处理 IP 被复用：旧设备的 IP 现在属于不同 MAC → 移除
	for mac, dev := range ctx.Devices {
		if newMAC, ok := ipToMAC[dev.IP]; ok && newMAC != mac {
			delete(ctx.Devices, mac)
			delete(ctx.Results, mac)
			delete(ctx.Printed, mac)
			fmt.Printf("  [-] IP 复用: %s 现属于 %s，移除旧设备 %s\n", dev.IP, newMAC, mac)
		}
	}

	// 2. 合并扫描结果
	added, updated := 0, 0
	for _, h := range hosts {
		mac := h.MAC
		ip := h.IP.String()
		if mac == "" {
			// 无法获取 MAC：降级用 IP 做 key，标记为临时
			mac = "tmp:" + ip
		}

		if dev, ok := ctx.Devices[mac]; ok {
			// 已存在：仅更新 IP（设备可能换了地址）
			if dev.IP != ip {
				dev.IP = ip
				dev.Host = h
				updated++
			}
		} else {
			// 新设备
			ctx.Devices[mac] = &Device{MAC: mac, IP: ip, Host: h}
			added++
		}
	}

	fmt.Printf("  [扫描] 设备 %d 台（新增 %d，IP更新 %d）\n", len(ctx.Devices), added, updated)
}

// testPending 对未测试或失败设备执行测试
func testPending(ctx *Context) {
	for mac, dev := range ctx.Devices {
		// 跳过: 已有结果且通过
		if rpt, ok := ctx.Results[mac]; ok && rpt != nil && rpt.Failed == 0 && rpt.Errors == 0 {
			continue
		}

		// 需要测试（结果为空 或 失败）
		fmt.Printf("  [测试] %s (%s) 开始压测...\n", dev.IP, dev.MAC)
		rpt := testDevice(ctx, dev.IP)
		ctx.Results[mac] = rpt

		// 保存测试报告
		saveDeviceReport(dev, rpt, ctx.Config.Report.Path)
	}
}

// testDevice 对单台设备执行 SSH 登录 + 压力测试
func testDevice(ctx *Context, ip string) *stress.Report {
	cfg := ctx.Config

	// SSH 连接
	client := sshclient.NewClient(
		cfg.SSH.User, cfg.SSH.Password, cfg.SSH.KeyFile,
		cfg.SSH.Port, cfg.Scan.Timeout,
	)
	conn, err := client.Connect(ip)
	if err != nil {
		fmt.Printf("  [测试] %s SSH 连接失败: %v\n", ip, err)
		return &stress.Report{
			IP:        ip,
			StartTime: time.Now(),
			Results: []stress.Result{{
				Type:   "ssh",
				Status: "fail",
				Error:  err.Error(),
			}},
			Failed: 1,
		}
	}
	defer conn.Close()

	// 执行压力测试
	runner := stress.NewRunner(
		stress.Config{
			Duration:   cfg.Stress.Duration,
			Threads:    cfg.Stress.Threads,
			DiskSizeMB: 1024,
			TempLogInt: 10 * time.Second,
		},
		mapStressTypes(cfg.Stress.Types),
		client,
	)
	return runner.Run(conn, ip)
}

// printStatus 打印状态变化（仅当首次或状态变化时）
func printStatus(ctx *Context) {
	for mac, rpt := range ctx.Results {
		dev := ctx.Devices[mac]
		label := devLabel(dev)
		status := statusOf(rpt)
		// 上次状态为空 或 与当前不一致 → 打印并记录
		if prev, ok := ctx.Printed[mac]; !ok || prev != status {
			icon := map[string]string{
				StatusPass: "✓",
				StatusFail: "✗",
			}[status]
			fmt.Printf("  [状态] %s %s → %s\n", label, icon, status)
			ctx.Printed[mac] = status
		}
	}

	// 未测试设备（无结果）: 首次打印一次 untested
	for mac, dev := range ctx.Devices {
		if _, hasResult := ctx.Results[mac]; !hasResult {
			if _, printed := ctx.Printed[mac]; !printed {
				fmt.Printf("  [状态] %s ○ → %s\n", devLabel(dev), StatusUntested)
				ctx.Printed[mac] = StatusUntested
			}
		}
	}
}

// devLabel 设备显示名: MAC (IP)
func devLabel(dev *Device) string {
	if dev == nil {
		return "?"
	}
	if strings.HasPrefix(dev.MAC, "tmp:") {
		return dev.IP + " (无MAC)"
	}
	return fmt.Sprintf("%s (%s)", dev.MAC, dev.IP)
}

// mapStressTypes 逗号分隔字符串 → TestType 列表
func mapStressTypes(s string) []stress.TestType {
	var types []stress.TestType
	for _, t := range strings.Split(s, ",") {
		t = strings.TrimSpace(t)
		types = append(types, stress.TestType(t))
	}
	return types
}

// statusOf 从测试报告得出状态
func statusOf(rpt *stress.Report) string {
	if rpt == nil {
		return StatusUntested
	}
	if rpt.Failed > 0 || rpt.Errors > 0 {
		return StatusFail
	}
	return StatusPass
}

// saveDeviceReport 保存单台设备的测试报告（文件名含 MAC 和 IP）
func saveDeviceReport(dev *Device, rpt *stress.Report, dir string) {
	if dir == "" {
		dir = "reports"
	}
	ts := time.Now().Format("20060102-150405")
	name := strings.ReplaceAll(dev.MAC, ":", "")
	path := fmt.Sprintf("%s/report-%s-%s-%s.txt", strings.TrimRight(dir, "/"), name, dev.IP, ts)

	r := report.NewBuilder(report.ConfigSummary{})
	r.StartPhase("压力测试")
	for _, res := range rpt.Results {
		r.AddStressResult(dev.IP, report.StressResult{
			Type:     string(res.Type),
			Status:   res.Status,
			Duration: res.Duration.String(),
			Error:    res.Error,
			Output:   res.Output,
		})
	}
	r.EndPhase(rpt.Summary())

	if err := r.Build().SaveFile(path, "text"); err != nil {
		fmt.Printf("  [报告] %s 保存失败: %v\n", dev.IP, err)
	} else {
		fmt.Printf("  [报告] %s 已保存: %s\n", dev.IP, path)
	}
}