zeaslity/cmii-uav-watchdog-project
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

123

Browse Source
This commit is contained in:
zeaslity
2025-03-13 11:22:44 +08:00
parent 34147b2f69
commit 4f8a8a6ff2
6 changed files with 117 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
cmii-uav-watchdog-otp/totp/totp.go
View File

@@ -77,6 +77,14 @@ type ValidateOpts struct {
Encoder otp.Encoder
}
func (opts *ValidateOpts) ConvertToValidateOpts(generateOpts GenerateOpts) {
opts.Period = generateOpts.Period
opts.Skew = 1
opts.Digits = generateOpts.Digits
opts.Algorithm = generateOpts.Algorithm
opts.Encoder = otp.EncoderDefault
}
// GenerateCodeCustom takes a timepoint and produces a passcode using a
// secret and the provided opts. (Under the hood, this is making an adapted
// call to hcmii-uav-watchdog-otp.GenerateCodeCustom)

2
cmii-uav-watchdog/main.go
View File

@@ -17,7 +17,7 @@ func main() {
return
}
// 初始化授权服务
// 初始化授权服务（使用单例模式）
authService := services.NewAuthService()
if authService == nil {
log.Fatalf("初始化授权服务失败")

74
cmii-uav-watchdog/services/auth_service.go
View File

@@ -12,6 +12,13 @@ import (
"time"
)
// 单例相关变量
var (
authServiceInstance *AuthService
authServiceOnce sync.Once
authServiceMutex sync.Mutex
)
// AuthService 授权服务
type AuthService struct {
mu sync.RWMutex
@@ -21,36 +28,51 @@ type AuthService struct {
initialized bool
}
// NewAuthService 创建授权服务
// NewAuthService 创建授权服务（单例模式）
func NewAuthService() *AuthService {
service := &AuthService{
hostInfoSet: make(map[string]models2.HostInfo),
totpService: NewTOTPService(),
initialized: false,
}
// 使用sync.Once确保初始化逻辑只执行一次
authServiceOnce.Do(func() {
authServiceMutex.Lock()
defer authServiceMutex.Unlock()
// 尝试从本地加载授权信息
service.loadAuthorizationInfo()
// 判断 项目级别的 TOTP密钥是否为空
// 若为空 则生成一个 二级TOTP密钥 然后持久化写入到授权文件中
if service.authorizationInfo.SecondTOTPSecret == "" {
secondTOTPSecret, err := service.totpService.GenerateTOTPSecret()
if err != nil {
log.Printf("生成二级TOTP密钥失败: %v", err)
return nil
// 如果实例已存在，直接返回
if authServiceInstance != nil {
return
}
service.authorizationInfo.SecondTOTPSecret = secondTOTPSecret
// 持久化写入到授权文件中
err = service.saveAuthorizationInfo()
if err != nil {
log.Printf("持久化写入授权文件失败: %v", err)
return nil
// 创建新实例
service := &AuthService{
hostInfoSet: make(map[string]models2.HostInfo),
totpService: NewTOTPService(),
initialized: false,
}
}
return service
// 尝试从本地加载授权信息
service.loadAuthorizationInfo()
// 判断 项目级别的 TOTP密钥是否为空
// 若为空 则生成一个 二级TOTP密钥 然后持久化写入到授权文件中
if service.authorizationInfo.SecondTOTPSecret == "" {
secondTOTPSecret, err := service.totpService.GenerateTierTwoTOTPSecret()
if err != nil {
log.Printf("生成二级TOTP密钥失败: %v", err)
return
}
service.authorizationInfo.SecondTOTPSecret = secondTOTPSecret
// 持久化写入到授权文件中
err = service.saveAuthorizationInfo()
if err != nil {
log.Printf("持久化写入授权文件失败: %v", err)
return
}
}
// 设置全局实例
authServiceInstance = service
})
return authServiceInstance
}
// AddHostInfo 添加主机信息
@@ -73,7 +95,7 @@ func (as *AuthService) GenerateAuthorizationFile() (*models2.AuthorizationFile,
}
// 生成TOTP验证码
totpCode, err := as.totpService.GenerateTOTP()
totpCode, err := as.totpService.GenerateTierTwoTOTPCode(as.authorizationInfo.SecondTOTPSecret)
if err != nil {
return nil, err
}
@@ -118,7 +140,7 @@ func (as *AuthService) ProcessAuthorizationCode(code models2.AuthorizationCode)
defer as.mu.Unlock()
// 验证TOTP
if err := as.totpService.VerifyTOTP(code.TOTPCode); err != nil {
if !as.totpService.VerifyTierTwoTOTPCode(code.TOTPCode, as.authorizationInfo.SecondTOTPSecret) {
return errors.New("无效的授权码: TOTP验证失败")
}

21
cmii-uav-watchdog/services/heartbeat_service.go
View File

@@ -3,6 +3,7 @@ package services
import (
"cmii-uav-watchdog-common/models"
"errors"
"log"
"time"
)
@@ -27,6 +28,11 @@ func (hs *HeartbeatService) ProcessHeartbeat(req models.HeartbeatRequest) (*mode
return nil, errors.New("无效的时间戳")
}
secondTOTPSecret := hs.authService.authorizationInfo.SecondTOTPSecret
if secondTOTPSecret == "" {
return nil, errors.New("二级TOTP密钥为空")
}
// 添加主机信息到集合
hs.authService.AddHostInfo(req.HostInfo)
@@ -36,17 +42,26 @@ func (hs *HeartbeatService) ProcessHeartbeat(req models.HeartbeatRequest) (*mode
Authorized: false,
TOTPCode: "",
Timestamp: time.Now().Unix(),
SecondTOTPSecret: "",
SecondTOTPSecret: secondTOTPSecret,
}, nil
}
// 检查totp密码是都有效
// 检查totp验证码是否有效
if !hs.totpService.VerifyTierTwoTOTPCode(req.TOTPCode, secondTOTPSecret) {
// 解析认证主机的相关信息
// 计算 请求时间与当前时间的时间差
diff := time.Now().Unix() - req.Timestamp
log.Printf("心跳请求时间与当前时间的时间差: %d", diff)
return nil, errors.New("无效的TOTP验证码，请检查系统时间是否正确！")
}
// 检查主机是否已授权
authorized := hs.authService.IsHostAuthorized(req.HostInfo)
// 生成TOTP验证码
totpCode, err := hs.totpService.GenerateTOTP()
totpCode, err := hs.totpService.GenerateTierTwoTOTPCode(secondTOTPSecret)
if err != nil {
return nil, err
}

59
cmii-uav-watchdog/services/totp_service.go
View File

@@ -2,7 +2,6 @@ package services
import (
"cmii-uav-watchdog/config"
"errors"
"log"
"time"
@@ -10,6 +9,17 @@ import (
"cmii-uav-watchdog-otp/totp"
)
var tierTwoTOTPSecretOpts = totp.GenerateOpts{
SecretSize: 32,
Issuer: "cmii-uav-watchdog",
AccountName: "cmii-uav-watchdog",
Period: 30,
Secret: []byte{},
Digits: otp.DigitsSix,
Algorithm: otp.AlgorithmSHA1,
Rand: nil,
}
// TOTPService TOTP服务
type TOTPService struct {
secret string
@@ -22,8 +32,8 @@ func NewTOTPService() *TOTPService {
}
}
// GenerateTOTP 生成TOTP验证码
func (ts *TOTPService) GenerateTOTP() (string, error) {
// GenerateTierOneTOTP 生成一级TOTP验证码
func (ts *TOTPService) GenerateTierOneTOTP() (string, error) {
// 使用当前时间生成TOTP
code, err := totp.GenerateCode(ts.secret, time.Now())
if err != nil {
@@ -33,29 +43,20 @@ func (ts *TOTPService) GenerateTOTP() (string, error) {
return code, nil
}
// VerifyTOTP 验证TOTP验证码
func (ts *TOTPService) VerifyTOTP(code string) error {
// VerifyTierOneTOTP 验证一级TOTP验证码
func (ts *TOTPService) VerifyTierOneTOTP(code string) bool {
// 验证TOTP
valid := totp.Validate(code, ts.secret)
if !valid {
return errors.New("无效的TOTP验证码")
return false
}
return nil
return true
}
// GenerateTOTPSecret 生成TOTP密钥
func (ts *TOTPService) GenerateTOTPSecret() (string, error) {
secret, err := totp.Generate(totp.GenerateOpts{
SecretSize: 32,
Issuer: "cmii-uav-watchdog",
AccountName: "cmii-uav-watchdog",
Period: 30,
Secret: []byte{},
Digits: otp.DigitsSix,
Algorithm: otp.AlgorithmSHA1,
Rand: nil,
})
// GenerateTierTwoTOTPSecret 生成二级TOTP密钥
func (ts *TOTPService) GenerateTierTwoTOTPSecret() (string, error) {
secret, err := totp.Generate(tierTwoTOTPSecretOpts)
if err != nil {
log.Printf("生成TOTP密钥失败: %v", err)
return "", err
@@ -63,3 +64,23 @@ func (ts *TOTPService) GenerateTOTPSecret() (string, error) {
return secret.Secret(), nil
}
// GenerateTierTwoTOTPCode 生成二级TOTP验证码
func (ts *TOTPService) GenerateTierTwoTOTPCode(secret string) (string, error) {
code, err := totp.GenerateCode(secret, time.Now())
if err != nil {
return "", err
}
return code, nil
}
// VerifyTierTwoTOTPCode 验证二级TOTP验证码
func (ts *TOTPService) VerifyTierTwoTOTPCode(code string, secret string) bool {
validateOpts := totp.ValidateOpts{}
validateOpts.ConvertToValidateOpts(tierTwoTOTPSecretOpts)
valid, err := totp.ValidateCustom(code, secret, time.Now(), validateOpts)
if err != nil {
return false
}
return valid
}

4
go.mod
View File

@@ -1,3 +1,3 @@
module cmii-uav-watchdog
module cmii-uav-watchdog-project
go 1.23
go 1.24