diff --git a/cmii-uav-watchdog-agent/cmd/watchdog-agent.go b/cmii-uav-watchdog-agent/cmd/watchdog-agent.go new file mode 100644 index 0000000..4ef89e4 --- /dev/null +++ b/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 +} diff --git a/cmii-uav-watchdog-agent/services/cpu_service.go b/cmii-uav-watchdog-agent/host_info/cpu_service.go similarity index 100% rename from cmii-uav-watchdog-agent/services/cpu_service.go rename to cmii-uav-watchdog-agent/host_info/cpu_service.go diff --git a/cmii-uav-watchdog-agent/services/disk_service.go b/cmii-uav-watchdog-agent/host_info/disk_service.go similarity index 100% rename from cmii-uav-watchdog-agent/services/disk_service.go rename to cmii-uav-watchdog-agent/host_info/disk_service.go diff --git a/cmii-uav-watchdog-agent/services/host_info.go b/cmii-uav-watchdog-agent/host_info/host_info.go similarity index 100% rename from cmii-uav-watchdog-agent/services/host_info.go rename to cmii-uav-watchdog-agent/host_info/host_info.go diff --git a/cmii-uav-watchdog-agent/services/mboard_service.go b/cmii-uav-watchdog-agent/host_info/mboard_service.go similarity index 100% rename from cmii-uav-watchdog-agent/services/mboard_service.go rename to cmii-uav-watchdog-agent/host_info/mboard_service.go diff --git a/cmii-uav-watchdog-agent/services/mem_service.go b/cmii-uav-watchdog-agent/host_info/mem_service.go similarity index 100% rename from cmii-uav-watchdog-agent/services/mem_service.go rename to cmii-uav-watchdog-agent/host_info/mem_service.go diff --git a/cmii-uav-watchdog-agent/services/net_service.go b/cmii-uav-watchdog-agent/host_info/net_service.go similarity index 100% rename from cmii-uav-watchdog-agent/services/net_service.go rename to cmii-uav-watchdog-agent/host_info/net_service.go diff --git a/cmii-uav-watchdog-agent/totp/auth.go b/cmii-uav-watchdog-agent/totp/auth.go new file mode 100644 index 0000000..205fc4f --- /dev/null +++ b/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 +} diff --git a/cmii-uav-watchdog-common/models/in_project_model.go b/cmii-uav-watchdog-common/models/in_project_model.go index 66b9ace..c8ff900 100644 --- a/cmii-uav-watchdog-common/models/in_project_model.go +++ b/cmii-uav-watchdog-common/models/in_project_model.go @@ -52,15 +52,6 @@ type HostInfo struct { MotherboardInfo MotherboardInfo `json:"motherboard_info"` } -//// HostInfo 主机信息模型 -//type HostInfo struct { -// UUID string `json:"uuid"` // 主机UUID -// CPU string `json:"cpu"` // CPU信息 -// Motherboard string `json:"motherboard"` // 主板信息 -// MAC string `json:"mac"` // MAC地址 -// Disk string `json:"disk"` // 硬盘信息 -//} - // HeartbeatRequest 心跳请求 type HeartbeatRequest struct { HostInfo HostInfo `json:"host_info"` // 主机信息