修改项目结构

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
+}

cmii-uav-watchdog-agent/totp/auth.go

@@ -0,0 +1,105 @@
+package totp
+
+import (
+	"crypto/hmac"
+	"crypto/sha1"
+	"encoding/base32"
+	"encoding/binary"
+	"fmt"
+	"strings"
+	"sync"
+	"time"
+)
+
+// TOTPConfig TOTP配置
+type TOTPConfig struct {
+	Secret    string        // TOTP密钥
+	Digits    int           // TOTP验证码长度
+	TimeStep  time.Duration // TOTP时间步长
+	Algorithm string        // TOTP算法
+}
+
+var (
+	defaultConfig = TOTPConfig{
+		Secret:    "",
+		Digits:    6,
+		TimeStep:  30 * time.Second,
+		Algorithm: "SHA1",
+	}
+	mu sync.RWMutex
+)
+
+// SetTOTPSecret 设置TOTP密钥
+func SetTOTPSecret(secret string) {
+	mu.Lock()
+	defer mu.Unlock()
+	defaultConfig.Secret = secret
+}
+
+// GetTOTPSecret 获取TOTP密钥
+func GetTOTPSecret() string {
+	mu.RLock()
+	defer mu.RUnlock()
+	return defaultConfig.Secret
+}
+
+// GenerateTOTPCode 生成TOTP验证码
+func GenerateTOTPCode() (string, error) {
+	mu.RLock()
+	config := defaultConfig
+	mu.RUnlock()
+
+	if config.Secret == "" {
+		return "", fmt.Errorf("TOTP密钥未设置")
+	}
+
+	// 确保密钥是Base32编码的
+	secret := strings.ToUpper(config.Secret)
+	secret = strings.ReplaceAll(secret, " ", "")
+	secretBytes, err := base32.StdEncoding.DecodeString(secret)
+	if err != nil {
+		return "", fmt.Errorf("解码TOTP密钥失败: %w", err)
+	}
+
+	// 获取当前时间戳并转换为TOTP时间计数器
+	timeCounter := uint64(time.Now().Unix()) / uint64(config.TimeStep.Seconds())
+
+	// 将时间计数器转换为字节数组
+	timeBytes := make([]byte, 8)
+	binary.BigEndian.PutUint64(timeBytes, timeCounter)
+
+	// 使用HMAC-SHA1计算TOTP值
+	h := hmac.New(sha1.New, secretBytes)
+	h.Write(timeBytes)
+	hash := h.Sum(nil)
+
+	// 根据RFC 6238，我们使用哈希的最后一个字节的低4位作为偏移
+	offset := hash[len(hash)-1] & 0x0f
+
+	// 从哈希中提取4字节并转换为整数
+	binary := binary.BigEndian.Uint32(hash[offset : offset+4])
+
+	// 屏蔽最高位并获取指定位数的数字
+	totp := binary & 0x7fffffff % uint32(pow10(config.Digits))
+
+	// 将数字格式化为字符串，填充前导零
+	return fmt.Sprintf("%0*d", config.Digits, totp), nil
+}
+
+// ValidateTOTPCode 验证TOTP验证码
+func ValidateTOTPCode(code string) bool {
+	generatedCode, err := GenerateTOTPCode()
+	if err != nil {
+		return false
+	}
+	return code == generatedCode
+}
+
+// pow10 计算10的n次方
+func pow10(n int) int {
+	result := 1
+	for i := 0; i < n; i++ {
+		result *= 10
+	}
+	return result
+}