修改项目结构

cmii-uav-watchdog-agent/totp/auth.go

@@ -0,0 +1,105 @@
+package totp
+
+import (
+	"crypto/hmac"
+	"crypto/sha1"
+	"encoding/base32"
+	"encoding/binary"
+	"fmt"
+	"strings"
+	"sync"
+	"time"
+)
+
+// TOTPConfig TOTP配置
+type TOTPConfig struct {
+	Secret    string        // TOTP密钥
+	Digits    int           // TOTP验证码长度
+	TimeStep  time.Duration // TOTP时间步长
+	Algorithm string        // TOTP算法
+}
+
+var (
+	defaultConfig = TOTPConfig{
+		Secret:    "",
+		Digits:    6,
+		TimeStep:  30 * time.Second,
+		Algorithm: "SHA1",
+	}
+	mu sync.RWMutex
+)
+
+// SetTOTPSecret 设置TOTP密钥
+func SetTOTPSecret(secret string) {
+	mu.Lock()
+	defer mu.Unlock()
+	defaultConfig.Secret = secret
+}
+
+// GetTOTPSecret 获取TOTP密钥
+func GetTOTPSecret() string {
+	mu.RLock()
+	defer mu.RUnlock()
+	return defaultConfig.Secret
+}
+
+// GenerateTOTPCode 生成TOTP验证码
+func GenerateTOTPCode() (string, error) {
+	mu.RLock()
+	config := defaultConfig
+	mu.RUnlock()
+
+	if config.Secret == "" {
+		return "", fmt.Errorf("TOTP密钥未设置")
+	}
+
+	// 确保密钥是Base32编码的
+	secret := strings.ToUpper(config.Secret)
+	secret = strings.ReplaceAll(secret, " ", "")
+	secretBytes, err := base32.StdEncoding.DecodeString(secret)
+	if err != nil {
+		return "", fmt.Errorf("解码TOTP密钥失败: %w", err)
+	}
+
+	// 获取当前时间戳并转换为TOTP时间计数器
+	timeCounter := uint64(time.Now().Unix()) / uint64(config.TimeStep.Seconds())
+
+	// 将时间计数器转换为字节数组
+	timeBytes := make([]byte, 8)
+	binary.BigEndian.PutUint64(timeBytes, timeCounter)
+
+	// 使用HMAC-SHA1计算TOTP值
+	h := hmac.New(sha1.New, secretBytes)
+	h.Write(timeBytes)
+	hash := h.Sum(nil)
+
+	// 根据RFC 6238，我们使用哈希的最后一个字节的低4位作为偏移
+	offset := hash[len(hash)-1] & 0x0f
+
+	// 从哈希中提取4字节并转换为整数
+	binary := binary.BigEndian.Uint32(hash[offset : offset+4])
+
+	// 屏蔽最高位并获取指定位数的数字
+	totp := binary & 0x7fffffff % uint32(pow10(config.Digits))
+
+	// 将数字格式化为字符串，填充前导零
+	return fmt.Sprintf("%0*d", config.Digits, totp), nil
+}
+
+// ValidateTOTPCode 验证TOTP验证码
+func ValidateTOTPCode(code string) bool {
+	generatedCode, err := GenerateTOTPCode()
+	if err != nil {
+		return false
+	}
+	return code == generatedCode
+}
+
+// pow10 计算10的n次方
+func pow10(n int) int {
+	result := 1
+	for i := 0; i < n; i++ {
+		result *= 10
+	}
+	return result
+}