102 lines
2.7 KiB
Go
102 lines
2.7 KiB
Go
package services
|
||
|
||
import (
|
||
"cmii-uav-watchdog-common/models"
|
||
"cmii-uav-watchdog-common/totp_tier_two"
|
||
"cmii-uav-watchdog-common/utils"
|
||
"cmii-uav-watchdog-common/wdd_log"
|
||
"errors"
|
||
"time"
|
||
)
|
||
|
||
// HeartbeatService 心跳服务
|
||
type HeartbeatService struct {
|
||
authService *AuthService
|
||
totpService *TOTPService
|
||
}
|
||
|
||
// NewHeartbeatService 创建心跳服务
|
||
func NewHeartbeatService() *HeartbeatService {
|
||
return &HeartbeatService{
|
||
authService: NewAuthService(),
|
||
totpService: NewTOTPService(),
|
||
}
|
||
}
|
||
|
||
// ProcessHeartbeat 处理心跳请求
|
||
func (hs *HeartbeatService) ProcessHeartbeat(req models.HeartbeatRequest) (*models.HeartbeatResponse, error) {
|
||
// 检查时间戳是否有效
|
||
if !hs.isTimestampValid(req.Timestamp) {
|
||
return nil, errors.New("无效的时间戳")
|
||
}
|
||
|
||
secondTOTPSecret := hs.authService.authorizationInfo.SecondTOTPSecret
|
||
if secondTOTPSecret == "" {
|
||
wdd_log.Error("二级TOTP密钥为空")
|
||
return nil, errors.New("二级TOTP密钥为空")
|
||
}
|
||
|
||
// 添加主机信息到集合
|
||
hs.authService.AddHostInfo(req.HostInfo)
|
||
|
||
// 检查是否携带totp，如果为空立即返回
|
||
if req.TOTPCode == "" {
|
||
return &models.HeartbeatResponse{
|
||
Authorized: false,
|
||
TOTPCode: "",
|
||
Timestamp: time.Now().Unix(),
|
||
SecondTOTPSecret: secondTOTPSecret,
|
||
}, nil
|
||
}
|
||
|
||
// 检查totp验证码是否有效
|
||
if !totp_tier_two.VerifyTierTwoTOTPCode(req.TOTPCode, secondTOTPSecret) {
|
||
// 解析认证主机的相关信息
|
||
|
||
// 计算 请求时间与当前时间的时间差
|
||
utils.CurentTimeUnix()
|
||
diff := time.Now().Unix() - req.Timestamp
|
||
wdd_log.Error("TOTP验证失败！心跳请求时间与当前时间的时间差: %d", diff)
|
||
|
||
return nil, errors.New("无效的TOTP验证码，请检查系统时间是否正确！")
|
||
}
|
||
|
||
// 检查主机是否已授权
|
||
authorized := hs.authService.IsHostAuthorized(req.HostInfo)
|
||
|
||
// 生成TOTP验证码
|
||
totpCode, err := totp_tier_two.GenerateTierTwoTOTPCode(secondTOTPSecret)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
// 返回心跳响应
|
||
return &models.HeartbeatResponse{
|
||
Authorized: authorized,
|
||
TOTPCode: totpCode,
|
||
Timestamp: time.Now().Unix(),
|
||
}, nil
|
||
}
|
||
|
||
// isTimestampValid 检查时间戳是否有效
|
||
func (hs *HeartbeatService) isTimestampValid(timestamp int64) bool {
|
||
now := time.Now().Unix()
|
||
diff := now - timestamp
|
||
|
||
// 允许5分钟的时间偏差
|
||
return diff < 300 && diff > -300
|
||
}
|
||
|
||
// GetAllHostInfo 获取所有主机信息
|
||
func (hs *HeartbeatService) GetAllHostInfo() (map[string]models.HostInfo, error) {
|
||
// 调用AuthService的方法获取所有主机信息
|
||
hostInfoMap := hs.authService.GetAllHostInfo()
|
||
|
||
// 如果没有主机信息，返回一个空的map
|
||
if len(hostInfoMap) == 0 {
|
||
return make(map[string]models.HostInfo), nil
|
||
}
|
||
|
||
return hostInfoMap, nil
|
||
}
|