添加配置加载和授权服务初始化，更新授权模型引用，重构授权文件生成和处理逻辑，优化 TOTP 密钥生成及错误处理

cmii-uav-watchdog/controllers/auth_controller.go

@@ -1,10 +1,11 @@
 package controllers
 
 import (
+	models2 "cmii-uav-watchdog-common/models"
 	"cmii-uav-watchdog/models"
 	"cmii-uav-watchdog/services"
 	"net/http"
-
+	
 	"github.com/gin-gonic/gin"
 )
 
@@ -32,7 +33,7 @@ func (ac *AuthController) GenerateAuthFile(c *gin.Context) {
 		})
 		return
 	}
-
+	
 	c.JSON(http.StatusOK, models.Response{
 		Code:    200,
 		Message: "生成授权文件成功",
@@ -42,7 +43,7 @@ func (ac *AuthController) GenerateAuthFile(c *gin.Context) {
 
 // ReceiveAuthCode 接收授权码
 func (ac *AuthController) ReceiveAuthCode(c *gin.Context) {
-	var authCode models.AuthorizationCode
+	var authCode models2.AuthorizationCode
 	if err := c.ShouldBindJSON(&authCode); err != nil {
 		c.JSON(http.StatusBadRequest, models.Response{
 			Code:    400,
@@ -51,7 +52,7 @@ func (ac *AuthController) ReceiveAuthCode(c *gin.Context) {
 		})
 		return
 	}
-
+	
 	// 处理授权码
 	err := ac.authService.ProcessAuthorizationCode(authCode)
 	if err != nil {
@@ -62,7 +63,7 @@ func (ac *AuthController) ReceiveAuthCode(c *gin.Context) {
 		})
 		return
 	}
-
+	
 	c.JSON(http.StatusOK, models.Response{
 		Code:    200,
 		Message: "处理授权码成功",
@@ -74,7 +75,7 @@ func (ac *AuthController) ReceiveAuthCode(c *gin.Context) {
 func (ac *AuthController) NotifyAuthInfo(c *gin.Context) {
 	// 获取授权信息
 	authInfo := ac.authService.GetAuthorizationInfo()
-
+	
 	c.JSON(http.StatusOK, models.Response{
 		Code:    200,
 		Message: "获取授权信息成功",

cmii-uav-watchdog/main.go

@@ -10,8 +10,19 @@ import (
 
 func main() {
 
+	// 初始化配置信息
+	err := config.LoadConfig("./config/config.yaml")
+	if err != nil {
+		log.Fatalf("加载配置文件失败: %v", err)
+		return
+	}
+
 	// 初始化授权服务
 	authService := services.NewAuthService()
+	if authService == nil {
+		log.Fatalf("初始化授权服务失败")
+		return
+	}
 
 	// 启动授权码检测定时任务
 	go func() {

cmii-uav-watchdog/services/auth_service.go

@@ -3,11 +3,9 @@ package services
 import (
 	models2 "cmii-uav-watchdog-common/models"
 	"cmii-uav-watchdog/config"
-	"cmii-uav-watchdog/models"
 	"cmii-uav-watchdog/utils"
 	"encoding/json"
 	"errors"
-	"io/ioutil"
 	"log"
 	"os"
 	"sync"
@@ -17,8 +15,8 @@ import (
 // AuthService 授权服务
 type AuthService struct {
 	mu                sync.RWMutex
-	hostInfoSet       map[string]models2.HostInfo // 主机信息集合
-	authorizationInfo models.AuthorizationStorage // 授权信息
+	hostInfoSet       map[string]models2.HostInfo  // 主机信息集合
+	authorizationInfo models2.AuthorizationStorage // 授权信息
 	totpService       *TOTPService
 	initialized       bool
 }
@@ -30,10 +28,28 @@ func NewAuthService() *AuthService {
 		totpService: NewTOTPService(),
 		initialized: false,
 	}
-	
+
 	// 尝试从本地加载授权信息
 	service.loadAuthorizationInfo()
-	
+
+	// 判断 项目级别的 TOTP密钥是否为空
+	// 若为空 则生成一个 二级TOTP密钥 然后持久化写入到授权文件中
+	if service.authorizationInfo.SecondTOTPSecret == "" {
+		secondTOTPSecret, err := service.totpService.GenerateTOTPSecret()
+		if err != nil {
+			log.Printf("生成二级TOTP密钥失败: %v", err)
+			return nil
+		}
+		service.authorizationInfo.SecondTOTPSecret = secondTOTPSecret
+
+		// 持久化写入到授权文件中
+		err = service.saveAuthorizationInfo()
+		if err != nil {
+			log.Printf("持久化写入授权文件失败: %v", err)
+			return nil
+		}
+	}
+
 	return service
 }
 
@@ -41,39 +57,39 @@ func NewAuthService() *AuthService {
 func (as *AuthService) AddHostInfo(hostInfo models2.HostInfo) {
 	as.mu.Lock()
 	defer as.mu.Unlock()
-	
+
 	hostKey := utils.GenerateHostKey(hostInfo)
 	as.hostInfoSet[hostKey] = hostInfo
 }
 
 // GenerateAuthorizationFile 生成授权文件
-func (as *AuthService) GenerateAuthorizationFile() (*models.AuthorizationFile, error) {
+func (as *AuthService) GenerateAuthorizationFile() (*models2.AuthorizationFile, error) {
 	as.mu.RLock()
 	defer as.mu.RUnlock()
-	
+
 	// 检查是否有主机信息
 	if len(as.hostInfoSet) == 0 {
 		return nil, errors.New("没有可用的主机信息")
 	}
-	
+
 	// 生成TOTP验证码
 	totpCode, err := as.totpService.GenerateTOTP()
 	if err != nil {
 		return nil, err
 	}
-	
+
 	// 获取当前时间
 	now := time.Now()
-	
+
 	// 准备初次授权时间
 	firstAuthTime := now
 	if as.initialized {
 		firstAuthTime = as.authorizationInfo.FirstAuthTime
 	}
-	
+
 	// 计算时间偏移
 	timeOffset := now.Unix() - firstAuthTime.Unix()
-	
+
 	// 加密主机信息
 	encryptedHosts := make([]string, 0)
 	for _, hostInfo := range as.hostInfoSet {
@@ -83,67 +99,67 @@ func (as *AuthService) GenerateAuthorizationFile() (*models.AuthorizationFile, e
 		}
 		encryptedHosts = append(encryptedHosts, encrypted)
 	}
-	
+
 	// 创建授权文件
-	authFile := &models.AuthorizationFile{
+	authFile := &models2.AuthorizationFile{
 		EncryptedHosts: encryptedHosts,
 		TOTPCode:       totpCode,
 		CurrentTime:    now,
 		FirstAuthTime:  firstAuthTime,
 		TimeOffset:     timeOffset,
 	}
-	
+
 	return authFile, nil
 }
 
 // ProcessAuthorizationCode 处理授权码
-func (as *AuthService) ProcessAuthorizationCode(code models.AuthorizationCode) error {
+func (as *AuthService) ProcessAuthorizationCode(code models2.AuthorizationCode) error {
 	as.mu.Lock()
 	defer as.mu.Unlock()
-	
+
 	// 验证TOTP
 	if err := as.totpService.VerifyTOTP(code.TOTPCode); err != nil {
 		return errors.New("无效的授权码: TOTP验证失败")
 	}
-	
+
 	// 获取当前时间
 	now := time.Now()
-	
+
 	// 准备初次授权时间
 	firstAuthTime := now
 	if as.initialized {
 		firstAuthTime = as.authorizationInfo.FirstAuthTime
 	}
-	
+
 	// 计算授权时间偏移
 	timeOffset := now.Unix() - code.CurrentTime.Unix()
-	
+
 	// 加密授权码
 	authCodeJSON, err := json.Marshal(code)
 	if err != nil {
 		return err
 	}
-	
+
 	encryptedCode, err := utils.Encrypt(string(authCodeJSON), config.GetConfig().Auth.Secret)
 	if err != nil {
 		return err
 	}
-	
+
 	// 保存授权信息
-	as.authorizationInfo = models.AuthorizationStorage{
+	as.authorizationInfo = models2.AuthorizationStorage{
 		EncryptedCode:   encryptedCode,
 		FirstAuthTime:   firstAuthTime,
 		TimeOffset:      timeOffset,
 		AuthorizedHosts: code.EncryptedHosts,
 	}
-	
+
 	// 保存到文件
 	if err := as.saveAuthorizationInfo(); err != nil {
 		return err
 	}
-	
+
 	as.initialized = true
-	
+
 	return nil
 }
 
@@ -151,24 +167,24 @@ func (as *AuthService) ProcessAuthorizationCode(code models.AuthorizationCode) e
 func (as *AuthService) IsHostAuthorized(hostInfo models2.HostInfo) bool {
 	as.mu.RLock()
 	defer as.mu.RUnlock()
-	
+
 	if !as.initialized {
 		return false
 	}
-	
+
 	// 加密主机信息
 	encrypted, err := utils.EncryptHostInfo(hostInfo)
 	if err != nil {
 		return false
 	}
-	
+
 	// 检查是否在已授权列表中
 	for _, host := range as.authorizationInfo.AuthorizedHosts {
 		if host == encrypted {
 			return true
 		}
 	}
-	
+
 	return false
 }
 
@@ -176,24 +192,24 @@ func (as *AuthService) IsHostAuthorized(hostInfo models2.HostInfo) bool {
 func (as *AuthService) VerifyAuthorizationTime() {
 	as.mu.RLock()
 	defer as.mu.RUnlock()
-	
+
 	if !as.initialized {
 		return
 	}
-	
+
 	// 获取当前时间和存储的初次授权时间
 	now := time.Now()
 	storedOffset := as.authorizationInfo.TimeOffset
-	
+
 	// 计算实际时间偏移
 	actualOffset := now.Unix() - as.authorizationInfo.FirstAuthTime.Unix()
-	
+
 	// 计算偏差
 	offsetDiff := actualOffset - storedOffset
-	
+
 	// 获取允许的时间偏移
 	allowedOffset := config.GetConfig().Auth.TimeOffsetAllowed
-	
+
 	// 检查偏差是否超过允许范围
 	if offsetDiff > allowedOffset || offsetDiff < -allowedOffset {
 		log.Printf("检测到时间篡改! 存储偏移: %d, 实际偏移: %d, 差值: %d",
@@ -206,14 +222,14 @@ func (as *AuthService) VerifyAuthorizationTime() {
 func (as *AuthService) GetAuthorizationInfo() interface{} {
 	as.mu.RLock()
 	defer as.mu.RUnlock()
-	
+
 	if !as.initialized {
 		return map[string]interface{}{
 			"authorized": false,
 			"message":    "未授权",
 		}
 	}
-	
+
 	return map[string]interface{}{
 		"authorized":       true,
 		"authorized_hosts": len(as.authorizationInfo.AuthorizedHosts),
@@ -227,31 +243,31 @@ func (as *AuthService) saveAuthorizationInfo() error {
 	if err != nil {
 		return err
 	}
-	
-	return ioutil.WriteFile(config.GetConfig().Auth.AuthFilePath, data, 0600)
+
+	return os.WriteFile(config.GetConfig().Auth.AuthFilePath, data, 0600)
 }
 
 // loadAuthorizationInfo 从文件加载授权信息
 func (as *AuthService) loadAuthorizationInfo() {
 	filePath := config.GetConfig().Auth.AuthFilePath
-	
+
 	// 检查文件是否存在
 	if _, err := os.Stat(filePath); os.IsNotExist(err) {
 		return
 	}
-	
-	data, err := ioutil.ReadFile(filePath)
+
+	data, err := os.ReadFile(filePath)
 	if err != nil {
 		log.Printf("读取授权文件失败: %v", err)
 		return
 	}
-	
-	var authInfo models.AuthorizationStorage
+
+	var authInfo models2.AuthorizationStorage
 	if err := json.Unmarshal(data, &authInfo); err != nil {
 		log.Printf("解析授权文件失败: %v", err)
 		return
 	}
-	
+
 	as.authorizationInfo = authInfo
 	as.initialized = true
 }

cmii-uav-watchdog/services/totp_service.go

@@ -3,8 +3,10 @@ package services
 import (
 	"cmii-uav-watchdog/config"
 	"errors"
+	"log"
 	"time"
 
+	otp "cmii-uav-watchdog-otp"
 	"cmii-uav-watchdog-otp/totp"
 )
 
@@ -41,3 +43,23 @@ func (ts *TOTPService) VerifyTOTP(code string) error {
 
 	return nil
 }
+
+// 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,
+	})
+	if err != nil {
+		log.Printf("生成TOTP密钥失败: %v", err)
+		return "", err
+	}
+
+	return secret.Secret(), nil
+}