sub2api/backend/internal/service/totp_service.go

package service

import (
	"context"
	"crypto/rand"
	"crypto/subtle"
	"encoding/hex"
	"fmt"
	"log/slog"
	"time"

	"github.com/pquerna/otp/totp"

	infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
)

var (
	ErrTotpNotEnabled      = infraerrors.BadRequest("TOTP_NOT_ENABLED", "totp feature is not enabled")
	ErrTotpAlreadyEnabled  = infraerrors.BadRequest("TOTP_ALREADY_ENABLED", "totp is already enabled for this account")
	ErrTotpNotSetup        = infraerrors.BadRequest("TOTP_NOT_SETUP", "totp is not set up for this account")
	ErrTotpInvalidCode     = infraerrors.BadRequest("TOTP_INVALID_CODE", "invalid totp code")
	ErrTotpSetupExpired    = infraerrors.BadRequest("TOTP_SETUP_EXPIRED", "totp setup session expired")
	ErrTotpTooManyAttempts = infraerrors.TooManyRequests("TOTP_TOO_MANY_ATTEMPTS", "too many verification attempts, please try again later")
	ErrVerifyCodeRequired  = infraerrors.BadRequest("VERIFY_CODE_REQUIRED", "email verification code is required")
	ErrPasswordRequired    = infraerrors.BadRequest("PASSWORD_REQUIRED", "password is required")
)

// TotpCache defines cache operations for TOTP service
type TotpCache interface {
	// Setup session methods
	GetSetupSession(ctx context.Context, userID int64) (*TotpSetupSession, error)
	SetSetupSession(ctx context.Context, userID int64, session *TotpSetupSession, ttl time.Duration) error
	DeleteSetupSession(ctx context.Context, userID int64) error

	// Login session methods (for 2FA login flow)
	GetLoginSession(ctx context.Context, tempToken string) (*TotpLoginSession, error)
	SetLoginSession(ctx context.Context, tempToken string, session *TotpLoginSession, ttl time.Duration) error
	DeleteLoginSession(ctx context.Context, tempToken string) error

	// Rate limiting
	IncrementVerifyAttempts(ctx context.Context, userID int64) (int, error)
	GetVerifyAttempts(ctx context.Context, userID int64) (int, error)
	ClearVerifyAttempts(ctx context.Context, userID int64) error

	// Step-up grant methods (敏感操作 sudo 窗口)
	SetStepUpGrant(ctx context.Context, userID int64, sessionKey string, ttl time.Duration) error
	HasStepUpGrant(ctx context.Context, userID int64, sessionKey string) (bool, error)
}

// SecretEncryptor defines encryption operations for TOTP secrets
type SecretEncryptor interface {
	Encrypt(plaintext string) (string, error)
	Decrypt(ciphertext string) (string, error)
}

// TotpSetupSession represents a TOTP setup session
type TotpSetupSession struct {
	Secret     string // Plain text TOTP secret (not encrypted yet)
	SetupToken string // Random token to verify setup request
	CreatedAt  time.Time
}

// TotpLoginSession represents a pending 2FA login session
type TotpLoginSession struct {
	UserID           int64
	Email            string
	TokenExpiry      time.Time
	PendingOAuthBind *PendingOAuthBindLoginSession `json:"pending_oauth_bind,omitempty"`
}

type PendingOAuthBindLoginSession struct {
	PendingSessionToken string `json:"pending_session_token,omitempty"`
	BrowserSessionKey   string `json:"browser_session_key,omitempty"`
}

// TotpStatus represents the TOTP status for a user
type TotpStatus struct {
	Enabled        bool       `json:"enabled"`
	EnabledAt      *time.Time `json:"enabled_at,omitempty"`
	FeatureEnabled bool       `json:"feature_enabled"`
}

// TotpSetupResponse represents the response for initiating TOTP setup
type TotpSetupResponse struct {
	Secret     string `json:"secret"`
	QRCodeURL  string `json:"qr_code_url"`
	SetupToken string `json:"setup_token"`
	Countdown  int    `json:"countdown"` // seconds until setup expires
}

const (
	totpSetupTTL    = 5 * time.Minute
	totpLoginTTL    = 5 * time.Minute
	totpAttemptsTTL = 15 * time.Minute
	maxTotpAttempts = 5
	totpIssuer      = "Sub2API"
)

// TotpService handles TOTP operations
type TotpService struct {
	userRepo          UserRepository
	encryptor         SecretEncryptor
	cache             TotpCache
	settingService    *SettingService
	emailService      *EmailService
	emailQueueService *EmailQueueService
}

// NewTotpService creates a new TOTP service
func NewTotpService(
	userRepo UserRepository,
	encryptor SecretEncryptor,
	cache TotpCache,
	settingService *SettingService,
	emailService *EmailService,
	emailQueueService *EmailQueueService,
) *TotpService {
	return &TotpService{
		userRepo:          userRepo,
		encryptor:         encryptor,
		cache:             cache,
		settingService:    settingService,
		emailService:      emailService,
		emailQueueService: emailQueueService,
	}
}

// GetStatus returns the TOTP status for a user
func (s *TotpService) GetStatus(ctx context.Context, userID int64) (*TotpStatus, error) {
	featureEnabled := s.settingService.IsTotpEnabled(ctx)

	user, err := s.userRepo.GetByID(ctx, userID)
	if err != nil {
		return nil, fmt.Errorf("get user: %w", err)
	}

	return &TotpStatus{
		Enabled:        user.TotpEnabled,
		EnabledAt:      user.TotpEnabledAt,
		FeatureEnabled: featureEnabled,
	}, nil
}

// usesEmailVerification 判断 TOTP 启用/停用时的身份校验方式。
// 管理员一律使用密码校验：管理员账号的邮箱常为占位地址收不到验证码，
// 且管理员凭证失守时攻击者往往同时控制通知邮箱，邮箱验证码不构成有效防线。
// 普通用户维持原有行为：开启邮箱验证时用邮箱验证码，否则用密码。
func (s *TotpService) usesEmailVerification(ctx context.Context, user *User) bool {
	return !user.IsAdmin() && s.settingService.IsEmailVerifyEnabled(ctx)
}

// verifyIdentity 按 usesEmailVerification 的结果校验邮箱验证码或密码。
func (s *TotpService) verifyIdentity(ctx context.Context, user *User, emailCode, password string) error {
	if s.usesEmailVerification(ctx, user) {
		if emailCode == "" {
			return ErrVerifyCodeRequired
		}
		return s.emailService.VerifyCode(ctx, user.Email, emailCode)
	}
	if password == "" {
		return ErrPasswordRequired
	}
	if !user.CheckPassword(password) {
		return ErrPasswordIncorrect
	}
	return nil
}

// InitiateSetup starts the TOTP setup process
// If email verification is enabled, emailCode is required; otherwise password is required
func (s *TotpService) InitiateSetup(ctx context.Context, userID int64, emailCode, password string) (*TotpSetupResponse, error) {
	// Check if TOTP feature is enabled globally
	if !s.settingService.IsTotpEnabled(ctx) {
		return nil, ErrTotpNotEnabled
	}

	// Get user and check if TOTP is already enabled
	user, err := s.userRepo.GetByID(ctx, userID)
	if err != nil {
		return nil, fmt.Errorf("get user: %w", err)
	}

	if user.TotpEnabled {
		return nil, ErrTotpAlreadyEnabled
	}

	if err := s.verifyIdentity(ctx, user, emailCode, password); err != nil {
		return nil, err
	}

	// Generate a new TOTP key
	key, err := totp.Generate(totp.GenerateOpts{
		Issuer:      totpIssuer,
		AccountName: user.Email,
	})
	if err != nil {
		return nil, fmt.Errorf("generate totp key: %w", err)
	}

	// Generate a random setup token
	setupToken, err := generateRandomToken(32)
	if err != nil {
		return nil, fmt.Errorf("generate setup token: %w", err)
	}

	// Store the setup session in cache
	session := &TotpSetupSession{
		Secret:     key.Secret(),
		SetupToken: setupToken,
		CreatedAt:  time.Now(),
	}

	if err := s.cache.SetSetupSession(ctx, userID, session, totpSetupTTL); err != nil {
		return nil, fmt.Errorf("store setup session: %w", err)
	}

	return &TotpSetupResponse{
		Secret:     key.Secret(),
		QRCodeURL:  key.URL(),
		SetupToken: setupToken,
		Countdown:  int(totpSetupTTL.Seconds()),
	}, nil
}

// CompleteSetup completes the TOTP setup by verifying the code
func (s *TotpService) CompleteSetup(ctx context.Context, userID int64, totpCode, setupToken string) error {
	// Check if TOTP feature is enabled globally
	if !s.settingService.IsTotpEnabled(ctx) {
		return ErrTotpNotEnabled
	}

	// Get the setup session
	session, err := s.cache.GetSetupSession(ctx, userID)
	if err != nil {
		return ErrTotpSetupExpired
	}

	if session == nil {
		return ErrTotpSetupExpired
	}

	// Verify the setup token (constant-time comparison)
	if subtle.ConstantTimeCompare([]byte(session.SetupToken), []byte(setupToken)) != 1 {
		return ErrTotpSetupExpired
	}

	// Verify the TOTP code
	if !totp.Validate(totpCode, session.Secret) {
		return ErrTotpInvalidCode
	}

	setupSecretPrefix := "N/A"
	if len(session.Secret) >= 4 {
		setupSecretPrefix = session.Secret[:4]
	}
	slog.Debug("totp_complete_setup_before_encrypt",
		"user_id", userID,
		"secret_len", len(session.Secret),
		"secret_prefix", setupSecretPrefix)

	// Encrypt the secret
	encryptedSecret, err := s.encryptor.Encrypt(session.Secret)
	if err != nil {
		return fmt.Errorf("encrypt totp secret: %w", err)
	}

	slog.Debug("totp_complete_setup_encrypted",
		"user_id", userID,
		"encrypted_len", len(encryptedSecret))

	// Verify encryption by decrypting
	decrypted, decErr := s.encryptor.Decrypt(encryptedSecret)
	if decErr != nil {
		slog.Debug("totp_complete_setup_verify_failed",
			"user_id", userID,
			"error", decErr)
	} else {
		decryptedPrefix := "N/A"
		if len(decrypted) >= 4 {
			decryptedPrefix = decrypted[:4]
		}
		slog.Debug("totp_complete_setup_verified",
			"user_id", userID,
			"original_len", len(session.Secret),
			"decrypted_len", len(decrypted),
			"match", session.Secret == decrypted,
			"decrypted_prefix", decryptedPrefix)
	}

	// Update user with encrypted TOTP secret
	if err := s.userRepo.UpdateTotpSecret(ctx, userID, &encryptedSecret); err != nil {
		return fmt.Errorf("update totp secret: %w", err)
	}

	// Enable TOTP for the user
	if err := s.userRepo.EnableTotp(ctx, userID); err != nil {
		return fmt.Errorf("enable totp: %w", err)
	}

	// Clean up the setup session
	_ = s.cache.DeleteSetupSession(ctx, userID)

	return nil
}

// Disable disables TOTP for a user
// If email verification is enabled, emailCode is required; otherwise password is required
func (s *TotpService) Disable(ctx context.Context, userID int64, emailCode, password string) error {
	// Get user
	user, err := s.userRepo.GetByID(ctx, userID)
	if err != nil {
		return fmt.Errorf("get user: %w", err)
	}

	if !user.TotpEnabled {
		return ErrTotpNotSetup
	}

	if err := s.verifyIdentity(ctx, user, emailCode, password); err != nil {
		return err
	}

	// Disable TOTP
	if err := s.userRepo.DisableTotp(ctx, userID); err != nil {
		return fmt.Errorf("disable totp: %w", err)
	}

	return nil
}

// VerifyCode verifies a TOTP code for a user
func (s *TotpService) VerifyCode(ctx context.Context, userID int64, code string) error {
	slog.Debug("totp_verify_code_called",
		"user_id", userID,
		"code_len", len(code))

	// Check rate limiting
	attempts, err := s.cache.GetVerifyAttempts(ctx, userID)
	if err == nil && attempts >= maxTotpAttempts {
		return ErrTotpTooManyAttempts
	}

	// Get user
	user, err := s.userRepo.GetByID(ctx, userID)
	if err != nil {
		slog.Debug("totp_verify_get_user_failed",
			"user_id", userID,
			"error", err)
		return infraerrors.InternalServer("TOTP_VERIFY_ERROR", "failed to verify totp code")
	}

	if !user.TotpEnabled || user.TotpSecretEncrypted == nil {
		slog.Debug("totp_verify_not_setup",
			"user_id", userID,
			"enabled", user.TotpEnabled,
			"has_secret", user.TotpSecretEncrypted != nil)
		return ErrTotpNotSetup
	}

	slog.Debug("totp_verify_encrypted_secret",
		"user_id", userID,
		"encrypted_len", len(*user.TotpSecretEncrypted))

	// Decrypt the secret
	secret, err := s.encryptor.Decrypt(*user.TotpSecretEncrypted)
	if err != nil {
		slog.Debug("totp_verify_decrypt_failed",
			"user_id", userID,
			"error", err)
		return infraerrors.InternalServer("TOTP_VERIFY_ERROR", "failed to verify totp code")
	}

	secretPrefix := "N/A"
	if len(secret) >= 4 {
		secretPrefix = secret[:4]
	}
	slog.Debug("totp_verify_decrypted",
		"user_id", userID,
		"secret_len", len(secret),
		"secret_prefix", secretPrefix)

	// Verify the code
	valid := totp.Validate(code, secret)
	slog.Debug("totp_verify_result",
		"user_id", userID,
		"valid", valid,
		"secret_len", len(secret),
		"secret_prefix", secretPrefix,
		"server_time", time.Now().UTC().Format(time.RFC3339))

	if !valid {
		// Increment failed attempts
		_, _ = s.cache.IncrementVerifyAttempts(ctx, userID)
		return ErrTotpInvalidCode
	}

	// Clear attempt counter on success
	_ = s.cache.ClearVerifyAttempts(ctx, userID)

	return nil
}

// StepUpGrantTTL 敏感操作 step-up 验证的有效窗口（sudo 模式）。
const StepUpGrantTTL = 15 * time.Minute

// VerifyStepUp 校验 TOTP 码并授予当前会话一段时间的 step-up 权限。
// 返回授权有效期，供前端展示/设置提醒。
func (s *TotpService) VerifyStepUp(ctx context.Context, userID int64, sessionKey, code string) (time.Duration, error) {
	if err := s.VerifyCode(ctx, userID, code); err != nil {
		return 0, err
	}
	if err := s.cache.SetStepUpGrant(ctx, userID, sessionKey, StepUpGrantTTL); err != nil {
		return 0, fmt.Errorf("store step-up grant: %w", err)
	}
	return StepUpGrantTTL, nil
}

// HasStepUpGrant 检查当前会话是否在 step-up 有效期内。
func (s *TotpService) HasStepUpGrant(ctx context.Context, userID int64, sessionKey string) (bool, error) {
	return s.cache.HasStepUpGrant(ctx, userID, sessionKey)
}

// CreateLoginSession creates a temporary login session for 2FA
func (s *TotpService) CreateLoginSession(ctx context.Context, userID int64, email string) (string, error) {
	return s.createLoginSession(ctx, userID, email, nil)
}

// CreatePendingOAuthBindLoginSession creates a temporary 2FA session that will
// finalize a pending OAuth bind after the TOTP code is verified.
func (s *TotpService) CreatePendingOAuthBindLoginSession(
	ctx context.Context,
	userID int64,
	email string,
	pendingSessionToken string,
	browserSessionKey string,
) (string, error) {
	return s.createLoginSession(ctx, userID, email, &PendingOAuthBindLoginSession{
		PendingSessionToken: pendingSessionToken,
		BrowserSessionKey:   browserSessionKey,
	})
}

func (s *TotpService) createLoginSession(
	ctx context.Context,
	userID int64,
	email string,
	pendingOAuthBind *PendingOAuthBindLoginSession,
) (string, error) {
	// Generate a random temp token
	tempToken, err := generateRandomToken(32)
	if err != nil {
		return "", fmt.Errorf("generate temp token: %w", err)
	}

	session := &TotpLoginSession{
		UserID:           userID,
		Email:            email,
		TokenExpiry:      time.Now().Add(totpLoginTTL),
		PendingOAuthBind: pendingOAuthBind,
	}

	if err := s.cache.SetLoginSession(ctx, tempToken, session, totpLoginTTL); err != nil {
		return "", fmt.Errorf("store login session: %w", err)
	}

	return tempToken, nil
}

// GetLoginSession retrieves a login session
func (s *TotpService) GetLoginSession(ctx context.Context, tempToken string) (*TotpLoginSession, error) {
	return s.cache.GetLoginSession(ctx, tempToken)
}

// DeleteLoginSession deletes a login session
func (s *TotpService) DeleteLoginSession(ctx context.Context, tempToken string) error {
	return s.cache.DeleteLoginSession(ctx, tempToken)
}

// IsTotpEnabledForUser checks if TOTP is enabled for a specific user
func (s *TotpService) IsTotpEnabledForUser(ctx context.Context, userID int64) (bool, error) {
	user, err := s.userRepo.GetByID(ctx, userID)
	if err != nil {
		return false, fmt.Errorf("get user: %w", err)
	}
	return user.TotpEnabled, nil
}

// MaskEmail masks an email address for display
func MaskEmail(email string) string {
	if len(email) < 3 {
		return "***"
	}

	atIdx := -1
	for i, c := range email {
		if c == '@' {
			atIdx = i
			break
		}
	}

	if atIdx == -1 || atIdx < 1 {
		return email[:1] + "***"
	}

	localPart := email[:atIdx]
	domain := email[atIdx:]

	if len(localPart) <= 2 {
		return localPart[:1] + "***" + domain
	}

	return localPart[:1] + "***" + localPart[len(localPart)-1:] + domain
}

// generateRandomToken generates a random hex-encoded token
func generateRandomToken(byteLength int) (string, error) {
	b := make([]byte, byteLength)
	if _, err := rand.Read(b); err != nil {
		return "", err
	}
	return hex.EncodeToString(b), nil
}

// VerificationMethod represents the method required for TOTP operations
type VerificationMethod struct {
	Method string `json:"method"` // "email" or "password"
}

// GetVerificationMethod returns the verification method for TOTP operations.
// 与 verifyIdentity 保持同一判定：管理员一律返回 password。
func (s *TotpService) GetVerificationMethod(ctx context.Context, userID int64) (*VerificationMethod, error) {
	user, err := s.userRepo.GetByID(ctx, userID)
	if err != nil {
		return nil, fmt.Errorf("get user: %w", err)
	}
	if s.usesEmailVerification(ctx, user) {
		return &VerificationMethod{Method: "email"}, nil
	}
	return &VerificationMethod{Method: "password"}, nil
}

// SendVerifyCode sends an email verification code for TOTP operations
func (s *TotpService) SendVerifyCode(ctx context.Context, userID int64, locale ...string) error {
	// Check if email verification is enabled
	if !s.settingService.IsEmailVerifyEnabled(ctx) {
		return infraerrors.BadRequest("EMAIL_VERIFY_NOT_ENABLED", "email verification is not enabled")
	}

	// Get user email
	user, err := s.userRepo.GetByID(ctx, userID)
	if err != nil {
		return fmt.Errorf("get user: %w", err)
	}

	// Get site name for email
	siteName := s.settingService.GetSiteName(ctx)

	// Send verification code via queue
	return s.emailQueueService.EnqueueVerifyCode(user.Email, siteName, firstEmailLocale(locale))
}