修改项目结构

cmii-uav-watchdog-agent/cmd/watchdog-agent.go

@@ -0,0 +1,148 @@
+package cmd
+
+import (
+	"cmii-uav-watchdog-agent/rpc"
+	"cmii-uav-watchdog-agent/services"
+	"cmii-uav-watchdog-agent/totp"
+	"cmii-uav-watchdog-common/models"
+	"fmt"
+	"log"
+	"os"
+	"os/signal"
+	"syscall"
+	"time"
+)
+
+const (
+	// 最大重试次数
+	maxRetryCount = 5
+
+	// 默认心跳检测间隔
+	defaultHeartbeatInterval = 30 * time.Second
+
+	// 检测失败后的等待间隔
+	failWaitInterval = 5 * time.Second
+
+	// 环境变量名称
+	appNameEnv = "APP_NAME"
+)
+
+// 启动心跳检测
+func StartHeartbeatDetection() {
+	log.Println("启动心跳检测任务...")
+
+	// 创建RPC客户端
+	client := rpc.NewClient(nil)
+
+	// 监听终止信号
+	signalChan := make(chan os.Signal, 1)
+	signal.Notify(signalChan, os.Interrupt, syscall.SIGTERM)
+
+	// 失败计数器
+	failCount := 0
+
+	// 心跳检测循环
+	for {
+		select {
+		case <-signalChan:
+			log.Println("收到终止信号，停止心跳检测")
+			return
+		default:
+			// 尝试发送心跳请求
+			authorized, err := sendHeartbeat(client)
+
+			if err != nil {
+				log.Printf("心跳检测失败: %v", err)
+				failCount++
+			} else if !authorized {
+				log.Println("未获得授权")
+				failCount++
+			} else {
+				// 检测成功，重置失败计数
+				failCount = 0
+				log.Println("心跳检测成功，已获得授权")
+			}
+
+			// 检查是否达到最大失败次数
+			if failCount >= maxRetryCount {
+				log.Printf("心跳检测连续失败 %d 次，发送终止信号", failCount)
+				// 发送终止信号给start_up.go
+				process, err := os.FindProcess(os.Getpid())
+				if err == nil {
+					process.Signal(syscall.SIGTERM)
+				}
+				return
+			}
+
+			// 等待下一次检测
+			if err != nil || !authorized {
+				// 失败后等待较短时间
+				time.Sleep(failWaitInterval)
+			} else {
+				// 成功后等待正常间隔
+				time.Sleep(defaultHeartbeatInterval)
+			}
+		}
+	}
+}
+
+// 发送心跳请求
+func sendHeartbeat(client *rpc.Client) (bool, error) {
+	// 1. 获取主机信息
+	hostInfoData := services.GetAllInfo()
+	hostInfo := models.HostInfo{
+		CPUInfo:         hostInfoData.CPUInfo,
+		DiskInfo:        hostInfoData.DiskInfo,
+		MemoryInfo:      hostInfoData.MemoryInfo,
+		NetInfo:         hostInfoData.NetInfo,
+		MotherboardInfo: hostInfoData.MotherboardInfo,
+	}
+
+	// 2. 获取应用名称
+	appName := os.Getenv(appNameEnv)
+	if appName == "" {
+		appName = "unknown-app"
+		log.Printf("警告: 环境变量 %s 未设置，使用默认值: %s", appNameEnv, appName)
+	}
+
+	// 构建心跳请求
+	request := &models.HeartbeatRequest{
+		HostInfo:  hostInfo,
+		Timestamp: time.Now().Unix(),
+		AppName:   appName,
+	}
+
+	// 3. 如果已有TOTP密钥，则生成TOTP验证码
+	totpSecret := totp.GetTOTPSecret()
+	if totpSecret != "" {
+		totpCode, err := totp.GenerateTOTPCode()
+		if err != nil {
+			log.Printf("生成TOTP验证码失败: %v", err)
+		} else {
+			request.TOTPCode = totpCode
+		}
+	}
+
+	// 4. 发送心跳请求
+	response, err := client.SendHeartbeatWithRetry(request, 10*time.Second)
+	if err != nil {
+		return false, fmt.Errorf("发送心跳请求失败: %w", err)
+	}
+
+	// 5. 处理响应
+	if response.SecondTOTPSecret != "" {
+		// 存储TOTP密钥
+		totp.SetTOTPSecret(response.SecondTOTPSecret)
+		log.Println("已更新TOTP密钥")
+	}
+
+	// 6. 如果有TOTP验证码，进行验证
+	if response.TOTPCode != "" && totpSecret != "" {
+		if !totp.ValidateTOTPCode(response.TOTPCode) {
+			log.Println("TOTP验证码验证失败")
+			return false, nil
+		}
+	}
+
+	return response.Authorized, nil
+}