user-system/internal/auth/totp.go

package auth

import (
	"bytes"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"crypto/subtle"
	"encoding/base32"
	"encoding/base64"
	"encoding/hex"
	"fmt"
	"image/png"
	"strings"
	"time"

	"github.com/pquerna/otp"
	"github.com/pquerna/otp/totp"
)

const (
	// TOTPIssuer 应用名称（显示在 Authenticator App 中）
	TOTPIssuer = "UserManagementSystem"
	// TOTPPeriod TOTP 时间步长（秒）
	TOTPPeriod = 30
	// TOTPDigits TOTP 位数
	TOTPDigits = 6
	// TOTPAlgorithm TOTP 算法（使用 SHA256 更安全）
	TOTPAlgorithm = otp.AlgorithmSHA256
	// RecoveryCodeCount 恢复码数量
	RecoveryCodeCount = 8
	// RecoveryCodeLength 每个恢复码的字节长度（生成后编码为 hex 字符串）
	RecoveryCodeLength = 5
)

// TOTPManager TOTP 管理器
type TOTPManager struct{}

// NewTOTPManager 创建 TOTP 管理器
func NewTOTPManager() *TOTPManager {
	return &TOTPManager{}
}

// TOTPSetup TOTP 初始化结果
type TOTPSetup struct {
	Secret        string   `json:"secret"`          // Base32 密钥（用户备用）
	QRCodeBase64  string   `json:"qr_code_base64"`  // Base64 编码的 PNG 二维码图片
	RecoveryCodes []string `json:"recovery_codes"`  // 一次性恢复码列表
}

// GenerateSecret 为指定用户生成 TOTP 密钥及二维码
func (m *TOTPManager) GenerateSecret(username string) (*TOTPSetup, error) {
	key, err := totp.Generate(totp.GenerateOpts{
		Issuer:      TOTPIssuer,
		AccountName: username,
		Period:      TOTPPeriod,
		Digits:      otp.DigitsSix,
		Algorithm:   TOTPAlgorithm,
	})
	if err != nil {
		return nil, fmt.Errorf("generate totp key failed: %w", err)
	}

	// 生成二维码图片
	img, err := key.Image(200, 200)
	if err != nil {
		return nil, fmt.Errorf("generate qr image failed: %w", err)
	}
	var buf bytes.Buffer
	if err := png.Encode(&buf, img); err != nil {
		return nil, fmt.Errorf("encode qr image failed: %w", err)
	}
	qrBase64 := base64.StdEncoding.EncodeToString(buf.Bytes())

	// 生成恢复码
	codes, err := generateRecoveryCodes(RecoveryCodeCount)
	if err != nil {
		return nil, fmt.Errorf("generate recovery codes failed: %w", err)
	}

	return &TOTPSetup{
		Secret:        key.Secret(),
		QRCodeBase64:  qrBase64,
		RecoveryCodes: codes,
	}, nil
}

// ValidateCode 验证用户输入的 TOTP 码（允许 ±1 个时间窗口的时钟偏差）
func (m *TOTPManager) ValidateCode(secret, code string) bool {
	// 注意：pquerna/otp 库的 ValidateCustom 与 GenerateCode 存在算法不匹配 bug（GenerateCode 固定用 SHA1）
	// 因此使用 totp.Validate() 代替，它内部正确处理算法检测
	return totp.Validate(strings.TrimSpace(code), secret)
}

// GenerateCurrentCode 生成当前时间的 TOTP 码（用于测试）
func (m *TOTPManager) GenerateCurrentCode(secret string) (string, error) {
	return totp.GenerateCode(secret, time.Now().UTC())
}

// ValidateRecoveryCode 验证恢复码（传入哈希后的已存储恢复码列表，返回匹配索引）
// 注意：调用方负责在验证后将该恢复码标记为已使用
// 使用恒定时间比较防止时序攻击
func ValidateRecoveryCode(inputCode string, storedCodes []string) (int, bool) {
	normalized := strings.ToUpper(strings.ReplaceAll(strings.TrimSpace(inputCode), "-", ""))
	for i, stored := range storedCodes {
		storedNormalized := strings.ToUpper(strings.ReplaceAll(stored, "-", ""))
		// 使用恒定时间比较防止时序攻击
		if subtle.ConstantTimeCompare([]byte(normalized), []byte(storedNormalized)) == 1 {
			return i, true
		}
	}
	return -1, false
}

// HashRecoveryCode 使用 SHA256 哈希恢复码（用于存储）
func HashRecoveryCode(code string) (string, error) {
	h := sha256.Sum256([]byte(code))
	return hex.EncodeToString(h[:]), nil
}

// VerifyRecoveryCode 验证恢复码（自动哈希后比较）
func VerifyRecoveryCode(inputCode string, hashedCodes []string) (int, bool) {
	hashedInput, err := HashRecoveryCode(inputCode)
	if err != nil {
		return -1, false
	}
	for i, hashed := range hashedCodes {
		if hmac.Equal([]byte(hashedInput), []byte(hashed)) {
			return i, true
		}
	}
	return -1, false
}

// generateRecoveryCodes 生成 N 个随机恢复码（格式：XXXXX-XXXXX）
func generateRecoveryCodes(count int) ([]string, error) {
	codes := make([]string, count)
	for i := 0; i < count; i++ {
		b := make([]byte, RecoveryCodeLength*2)
		if _, err := rand.Read(b); err != nil {
			return nil, err
		}
		encoded := base32.StdEncoding.EncodeToString(b)
		// 格式化为 XXXXX-XXXXX
		part := strings.ToUpper(encoded[:10])
		codes[i] = part[:5] + "-" + part[5:]
	}
	return codes, nil
}