版本封存

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

@@ -1,76 +1,95 @@
-package cmd
+package main
 
 import (
 	"cmii-uav-watchdog-agent/host_info"
 	"cmii-uav-watchdog-agent/rpc"
-	"cmii-uav-watchdog-agent/totp"
 	"cmii-uav-watchdog-common/models"
-	"fmt"
-	"log"
+	"cmii-uav-watchdog-common/totp_tier_two"
+	"cmii-uav-watchdog-common/wdd_log"
 	"os"
-	"os/signal"
 	"syscall"
 	"time"
 )
 
-const (
+var (
 	// 最大重试次数
-	maxRetryCount = 5
+	maxRetryCount = 12
 
 	// 默认心跳检测间隔
-	defaultHeartbeatInterval = 30 * time.Second
+	defaultHeartbeatInterval = 2 * time.Hour
 
 	// 检测失败后的等待间隔
-	failWaitInterval = 5 * time.Second
-
-	// 环境变量名称
-	appNameEnv = "APP_NAME"
+	failWaitInterval = 1 * time.Hour
 )
 
-// 启动心跳检测
-func StartHeartbeatDetection() {
-	log.Println("启动心跳检测任务...")
+var tierTwoTotpSecret = ""
+
+// StartHeartbeatDetection 启动心跳检测
+func StartHeartbeatDetection(signalChan chan os.Signal) {
+	wdd_log.Info("启动心跳检测任务...")
+
+	// variable
+	var err error
+
+	// 如果Debug模式，那么使用环境变量赋值
+	if DebugMode {
+		// 如果从在环境变量 那么使用环境变量赋值
+		heartbeatInterval := os.Getenv("WATCHDOG_AGENT_HEARTBEAT_INTERVAL")
+		if heartbeatInterval != "" {
+			defaultHeartbeatInterval, err = time.ParseDuration(heartbeatInterval)
+			if err != nil {
+				wdd_log.Error("无法解析环境变量: %v", err)
+			}
+			wdd_log.Info("已更新心跳检测间隔 => %s", defaultHeartbeatInterval)
+		}
+
+		failWaitIntervalEnv := os.Getenv("WATCHDOG_AGENT_FAIL_WAIT_INTERVAL")
+		if failWaitIntervalEnv != "" {
+			failWaitInterval, err = time.ParseDuration(failWaitIntervalEnv)
+			if err != nil {
+				wdd_log.Error("无法解析环境变量: %v", err)
+			}
+			wdd_log.Info("已更新心跳检测失败等待间隔 => %s", failWaitInterval)
+		}
+
+		//
+	}
 
 	// 创建RPC客户端
-	client := rpc.NewClient(nil)
+	heartbeatURL := os.Getenv("WATCHDOG_HEARTBEAT_URL")
+	client := rpc.NewClient(nil, heartbeatURL)
 
-	// 监听终止信号
-	signalChan := make(chan os.Signal, 1)
-	signal.Notify(signalChan, os.Interrupt, syscall.SIGTERM)
+	wdd_log.Info("心跳检测URL: %s", client.GetHeartbeatURL())
 
 	// 失败计数器
-	failCount := 0
+	failCount := 1
 
 	// 心跳检测循环
 	for {
 		select {
 		case <-signalChan:
-			log.Println("收到终止信号，停止心跳检测")
+			wdd_log.Info("收到终止信号，停止心跳检测")
 			return
 		default:
 			// 尝试发送心跳请求
 			authorized, err := sendHeartbeat(client)
 
 			if err != nil {
-				log.Printf("心跳检测失败: %v", err)
+				wdd_log.Error("第 %d 次心跳检测失败: %v", failCount, err)
 				failCount++
 			} else if !authorized {
-				log.Println("未获得授权")
+				wdd_log.Warn("第 %d 次心跳检测未获得授权", failCount)
 				failCount++
 			} else {
 				// 检测成功，重置失败计数
-				failCount = 0
-				log.Println("心跳检测成功，已获得授权")
+				failCount = 1
+				wdd_log.Info("第 %d 次心跳检测成功，已获得授权", failCount)
 			}
 
 			// 检查是否达到最大失败次数
 			if failCount >= maxRetryCount {
-				log.Printf("心跳检测连续失败 %d 次，发送终止信号", failCount)
-				// 发送终止信号给start_up.go
-				process, err := os.FindProcess(os.Getpid())
-				if err == nil {
-					process.Signal(syscall.SIGTERM)
-				}
+				wdd_log.Fatal("心跳检测连续失败 %d 次，发送终止信号！！", failCount)
+				signalChan <- syscall.SIGTERM
 				return
 			}
 
@@ -88,8 +107,9 @@ func StartHeartbeatDetection() {
 
 // 发送心跳请求
 func sendHeartbeat(client *rpc.Client) (bool, error) {
+
 	// 1. 获取主机信息
-	hostInfoData := services.GetAllInfo()
+	hostInfoData := host_info.GetAllInfo()
 	hostInfo := models.HostInfo{
 		SystemInfo: hostInfoData.SystemInfo,
 		CPUInfo:    hostInfoData.CPUInfo,
@@ -98,26 +118,21 @@ func sendHeartbeat(client *rpc.Client) (bool, error) {
 		NetInfo:    hostInfoData.NetInfo,
 	}
 
-	// 2. 获取应用名称
-	appName := os.Getenv(appNameEnv)
-	if appName == "" {
-		appName = "unknown-app"
-		log.Printf("警告: 环境变量 %s 未设置，使用默认值: %s", appNameEnv, appName)
-	}
+	// 获取环境信息
+	envInfo := GetEnvInfo()
 
 	// 构建心跳请求
 	request := &models.HeartbeatRequest{
 		HostInfo:  hostInfo,
 		Timestamp: time.Now().Unix(),
-		AppName:   appName,
+		EnvInfo:   envInfo,
 	}
 
 	// 3. 如果已有TOTP密钥，则生成TOTP验证码
-	totpSecret := totp.GetTOTPSecret()
-	if totpSecret != "" {
-		totpCode, err := totp.GenerateTOTPCode()
+	if tierTwoTotpSecret != "" {
+		totpCode, err := totp_tier_two.GenerateTierTwoTOTPCode(tierTwoTotpSecret)
 		if err != nil {
-			log.Printf("生成TOTP验证码失败: %v", err)
+			wdd_log.Error("生成TOTP验证码失败: %v", err)
 		} else {
 			request.TOTPCode = totpCode
 		}
@@ -126,20 +141,24 @@ func sendHeartbeat(client *rpc.Client) (bool, error) {
 	// 4. 发送心跳请求
 	response, err := client.SendHeartbeatWithRetry(request, 10*time.Second)
 	if err != nil {
-		return false, fmt.Errorf("发送心跳请求失败: %w", err)
+		return false, err
+	}
+
+	if response == nil {
+		return false, err
 	}
 
 	// 5. 处理响应
 	if response.SecondTOTPSecret != "" {
 		// 存储TOTP密钥
-		totp.SetTOTPSecret(response.SecondTOTPSecret)
-		log.Println("已更新TOTP密钥")
+		tierTwoTotpSecret = response.SecondTOTPSecret
+		wdd_log.Info("已更新TOTP密钥 => %s", tierTwoTotpSecret)
 	}
 
 	// 6. 如果有TOTP验证码，进行验证
-	if response.TOTPCode != "" && totpSecret != "" {
-		if !totp.ValidateTOTPCode(response.TOTPCode) {
-			log.Println("TOTP验证码验证失败")
+	if response.TOTPCode != "" && tierTwoTotpSecret != "" {
+		if !totp_tier_two.VerifyTierTwoTOTPCode(response.TOTPCode, tierTwoTotpSecret) {
+			wdd_log.Warn("TOTP验证码验证失败")
 			return false, nil
 		}
 	}