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

This commit is contained in:
zeaslity
2025-03-13 10:36:23 +08:00
parent 25dbc87a73
commit 34147b2f69
5 changed files with 133 additions and 55 deletions

View File

@@ -0,0 +1,28 @@
package models
import "time"
// AuthorizationFile 授权文件模型
type AuthorizationFile struct {
EncryptedHosts []string `json:"encrypted_hosts"` // 加密后的主机信息列表
TOTPCode string `json:"totp_code"` // TOTP验证码
CurrentTime time.Time `json:"current_time"` // 当前系统时间
FirstAuthTime time.Time `json:"first_auth_time"` // 初次授权时间
TimeOffset int64 `json:"time_offset"` // 授权时间偏移
}
// AuthorizationCode 授权码模型
type AuthorizationCode struct {
TOTPCode string `json:"totp_code"` // TOTP验证码
CurrentTime time.Time `json:"current_time"` // 当前系统时间
EncryptedHosts []string `json:"encrypted_hosts"` // 授权主机的加密字符串列表
}
// AuthorizationStorage 授权存储信息
type AuthorizationStorage struct {
EncryptedCode string `json:"encrypted_code"` // 加密后的授权码
FirstAuthTime time.Time `json:"first_auth_time"` // 初次授权时间
TimeOffset int64 `json:"time_offset"` // 授权时间偏移
AuthorizedHosts []string `json:"authorized_hosts"` // 已授权主机列表
SecondTOTPSecret string `json:"second_totp_secret"` // 第二级的totp密钥
}

View File

@@ -1,6 +1,7 @@
package controllers
import (
models2 "cmii-uav-watchdog-common/models"
"cmii-uav-watchdog/models"
"cmii-uav-watchdog/services"
"net/http"
@@ -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,

View File

@@ -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() {

View File

@@ -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
}
@@ -34,6 +32,24 @@ func NewAuthService() *AuthService {
// 尝试从本地加载授权信息
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
}
@@ -47,7 +63,7 @@ func (as *AuthService) AddHostInfo(hostInfo models2.HostInfo) {
}
// GenerateAuthorizationFile 生成授权文件
func (as *AuthService) GenerateAuthorizationFile() (*models.AuthorizationFile, error) {
func (as *AuthService) GenerateAuthorizationFile() (*models2.AuthorizationFile, error) {
as.mu.RLock()
defer as.mu.RUnlock()
@@ -85,7 +101,7 @@ func (as *AuthService) GenerateAuthorizationFile() (*models.AuthorizationFile, e
}
// 创建授权文件
authFile := &models.AuthorizationFile{
authFile := &models2.AuthorizationFile{
EncryptedHosts: encryptedHosts,
TOTPCode: totpCode,
CurrentTime: now,
@@ -97,7 +113,7 @@ func (as *AuthService) GenerateAuthorizationFile() (*models.AuthorizationFile, e
}
// ProcessAuthorizationCode 处理授权码
func (as *AuthService) ProcessAuthorizationCode(code models.AuthorizationCode) error {
func (as *AuthService) ProcessAuthorizationCode(code models2.AuthorizationCode) error {
as.mu.Lock()
defer as.mu.Unlock()
@@ -130,7 +146,7 @@ func (as *AuthService) ProcessAuthorizationCode(code models.AuthorizationCode) e
}
// 保存授权信息
as.authorizationInfo = models.AuthorizationStorage{
as.authorizationInfo = models2.AuthorizationStorage{
EncryptedCode: encryptedCode,
FirstAuthTime: firstAuthTime,
TimeOffset: timeOffset,
@@ -228,7 +244,7 @@ func (as *AuthService) saveAuthorizationInfo() error {
return err
}
return ioutil.WriteFile(config.GetConfig().Auth.AuthFilePath, data, 0600)
return os.WriteFile(config.GetConfig().Auth.AuthFilePath, data, 0600)
}
// loadAuthorizationInfo 从文件加载授权信息
@@ -240,13 +256,13 @@ func (as *AuthService) loadAuthorizationInfo() {
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

View File

@@ -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
}