cmii-uav-watchdog-project/cmii-uav-watchdog/services/auth_service.go

package services

import (
	models2 "cmii-uav-watchdog-common/models"
	"cmii-uav-watchdog/config"
	"cmii-uav-watchdog/utils"
	"encoding/json"
	"errors"
	"log"
	"os"
	"sync"
	"time"
)

// AuthService 授权服务
type AuthService struct {
	mu                sync.RWMutex
	hostInfoSet       map[string]models2.HostInfo  // 主机信息集合
	authorizationInfo models2.AuthorizationStorage // 授权信息
	totpService       *TOTPService
	initialized       bool
}

// NewAuthService 创建授权服务
func NewAuthService() *AuthService {
	service := &AuthService{
		hostInfoSet: make(map[string]models2.HostInfo),
		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
}

// AddHostInfo 添加主机信息
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() (*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 {
		encrypted, err := utils.EncryptHostInfo(hostInfo)
		if err != nil {
			return nil, err
		}
		encryptedHosts = append(encryptedHosts, encrypted)
	}

	// 创建授权文件
	authFile := &models2.AuthorizationFile{
		EncryptedHosts: encryptedHosts,
		TOTPCode:       totpCode,
		CurrentTime:    now,
		FirstAuthTime:  firstAuthTime,
		TimeOffset:     timeOffset,
	}

	return authFile, nil
}

// ProcessAuthorizationCode 处理授权码
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 = models2.AuthorizationStorage{
		EncryptedCode:   encryptedCode,
		FirstAuthTime:   firstAuthTime,
		TimeOffset:      timeOffset,
		AuthorizedHosts: code.EncryptedHosts,
	}

	// 保存到文件
	if err := as.saveAuthorizationInfo(); err != nil {
		return err
	}

	as.initialized = true

	return nil
}

// IsHostAuthorized 检查主机是否已授权
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
}

// VerifyAuthorizationTime 验证授权时间有效性
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",
			storedOffset, actualOffset, offsetDiff)
		// 这里可以添加更多的处理逻辑，例如发送警告、禁用授权等
	}
}

// GetAuthorizationInfo 获取授权信息
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),
		"first_auth_time":  as.authorizationInfo.FirstAuthTime,
	}
}

// saveAuthorizationInfo 保存授权信息到文件
func (as *AuthService) saveAuthorizationInfo() error {
	data, err := json.Marshal(as.authorizationInfo)
	if err != nil {
		return err
	}

	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 := os.ReadFile(filePath)
	if err != nil {
		log.Printf("读取授权文件失败: %v", err)
		return
	}

	var authInfo models2.AuthorizationStorage
	if err := json.Unmarshal(data, &authInfo); err != nil {
		log.Printf("解析授权文件失败: %v", err)
		return
	}

	as.authorizationInfo = authInfo
	as.initialized = true
}