123

cmii-uav-watchdog/services/totp_service.go

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