user-system/internal/service/totp.go

package service

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"

	"github.com/user-management-system/internal/auth"
)

// TOTPService manages 2FA setup, enable/disable, and verification.
type TOTPService struct {
	userRepo    userRepositoryInterface
	totpManager *auth.TOTPManager
}

func NewTOTPService(userRepo userRepositoryInterface) *TOTPService {
	return &TOTPService{
		userRepo:    userRepo,
		totpManager: auth.NewTOTPManager(),
	}
}

type SetupTOTPResponse struct {
	Secret        string   `json:"secret"`
	QRCodeBase64  string   `json:"qr_code_base64"`
	RecoveryCodes []string `json:"recovery_codes"`
}

func (s *TOTPService) SetupTOTP(ctx context.Context, userID int64) (*SetupTOTPResponse, error) {
	user, err := s.userRepo.GetByID(ctx, userID)
	if err != nil {
		return nil, fmt.Errorf("\u7528\u6237\u4e0d\u5b58\u5728")
	}
	if user.TOTPEnabled {
		return nil, errors.New("2FA \u5df2\u7ecf\u542f\u7528\uff0c\u5982\u9700\u91cd\u7f6e\u8bf7\u5148\u7981\u7528")
	}

	setup, err := s.totpManager.GenerateSecret(user.Username)
	if err != nil {
		return nil, fmt.Errorf("\u751f\u6210 TOTP \u5bc6\u94a5\u5931\u8d25: %w", err)
	}

	// Persist the generated secret and recovery codes before activation.
	user.TOTPSecret = setup.Secret
	// Hash recovery codes before storing (SEC-03 fix)
	hashedCodes := make([]string, len(setup.RecoveryCodes))
	for i, code := range setup.RecoveryCodes {
		hashedCodes[i], _ = auth.HashRecoveryCode(code)
	}
	codesJSON, _ := json.Marshal(hashedCodes)
	user.TOTPRecoveryCodes = string(codesJSON)

	if err := s.userRepo.UpdateTOTP(ctx, user); err != nil {
		return nil, fmt.Errorf("\u4fdd\u5b58 TOTP \u5bc6\u94a5\u5931\u8d25: %w", err)
	}

	return &SetupTOTPResponse{
		Secret:        setup.Secret,
		QRCodeBase64:  setup.QRCodeBase64,
		RecoveryCodes: setup.RecoveryCodes,
	}, nil
}

func (s *TOTPService) EnableTOTP(ctx context.Context, userID int64, code string) error {
	user, err := s.userRepo.GetByID(ctx, userID)
	if err != nil {
		return fmt.Errorf("\u7528\u6237\u4e0d\u5b58\u5728")
	}
	if user.TOTPSecret == "" {
		return errors.New("\u8bf7\u5148\u521d\u59cb\u5316 2FA\uff0c\u83b7\u53d6\u4e8c\u7ef4\u7801\u540e\u518d\u6fc0\u6d3b")
	}
	if user.TOTPEnabled {
		return errors.New("2FA \u5df2\u542f\u7528")
	}

	if !s.totpManager.ValidateCode(user.TOTPSecret, code) {
		return errors.New("\u9a8c\u8bc1\u7801\u9519\u8bef\u6216\u5df2\u8fc7\u671f")
	}

	user.TOTPEnabled = true
	return s.userRepo.UpdateTOTP(ctx, user)
}

func (s *TOTPService) DisableTOTP(ctx context.Context, userID int64, code string) error {
	user, err := s.userRepo.GetByID(ctx, userID)
	if err != nil {
		return fmt.Errorf("\u7528\u6237\u4e0d\u5b58\u5728")
	}
	if !user.TOTPEnabled {
		return errors.New("2FA \u672a\u542f\u7528")
	}

	valid := s.totpManager.ValidateCode(user.TOTPSecret, code)
	if !valid {
		var hashedCodes []string
		if user.TOTPRecoveryCodes != "" {
			_ = json.Unmarshal([]byte(user.TOTPRecoveryCodes), &hashedCodes)
		}
		_, matched := auth.VerifyRecoveryCode(code, hashedCodes)
		if !matched {
			return errors.New("\u9a8c\u8bc1\u7801\u6216\u6062\u590d\u7801\u9519\u8bef")
		}
	}

	user.TOTPEnabled = false
	user.TOTPSecret = ""
	user.TOTPRecoveryCodes = ""
	return s.userRepo.UpdateTOTP(ctx, user)
}

func (s *TOTPService) VerifyTOTP(ctx context.Context, userID int64, code string) error {
	user, err := s.userRepo.GetByID(ctx, userID)
	if err != nil {
		return fmt.Errorf("\u7528\u6237\u4e0d\u5b58\u5728")
	}
	if !user.TOTPEnabled {
		return nil
	}

	if s.totpManager.ValidateCode(user.TOTPSecret, code) {
		return nil
	}

	var storedCodes []string
	if user.TOTPRecoveryCodes != "" {
		_ = json.Unmarshal([]byte(user.TOTPRecoveryCodes), &storedCodes)
	}
	idx, matched := auth.ValidateRecoveryCode(code, storedCodes)
	if !matched {
		return errors.New("\u9a8c\u8bc1\u7801\u9519\u8bef\u6216\u5df2\u8fc7\u671f")
	}

	storedCodes = append(storedCodes[:idx], storedCodes[idx+1:]...)
	codesJSON, _ := json.Marshal(storedCodes)
	user.TOTPRecoveryCodes = string(codesJSON)
	_ = s.userRepo.UpdateTOTP(ctx, user)
	return nil
}

func (s *TOTPService) GetTOTPStatus(ctx context.Context, userID int64) (bool, error) {
	user, err := s.userRepo.GetByID(ctx, userID)
	if err != nil {
		return false, fmt.Errorf("\u7528\u6237\u4e0d\u5b58\u5728")
	}
	return user.TOTPEnabled, nil
}