149 lines
3.4 KiB
Go
149 lines
3.4 KiB
Go
package cmd
|
||
|
||
import (
|
||
"cmii-uav-watchdog-agent/host_info"
|
||
"cmii-uav-watchdog-agent/rpc"
|
||
"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{
|
||
SystemInfo: hostInfoData.SystemInfo,
|
||
CPUInfo: hostInfoData.CPUInfo,
|
||
DiskInfo: hostInfoData.DiskInfo,
|
||
MemoryInfo: hostInfoData.MemoryInfo,
|
||
NetInfo: hostInfoData.NetInfo,
|
||
}
|
||
|
||
// 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
|
||
}
|