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 } // 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 } // 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 } // 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 } // Verify identity based on email verification setting if s.settingService.IsEmailVerifyEnabled(ctx) { // Email verification enabled - verify email code if emailCode == "" { return nil, ErrVerifyCodeRequired } if err := s.emailService.VerifyCode(ctx, user.Email, emailCode); err != nil { return nil, err } } else { // Email verification disabled - verify password if password == "" { return nil, ErrPasswordRequired } if !user.CheckPassword(password) { return nil, ErrPasswordIncorrect } } // 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 } // Verify identity based on email verification setting if s.settingService.IsEmailVerifyEnabled(ctx) { // Email verification enabled - verify email code if emailCode == "" { return ErrVerifyCodeRequired } if err := s.emailService.VerifyCode(ctx, user.Email, emailCode); err != nil { return err } } else { // Email verification disabled - verify password if password == "" { return ErrPasswordRequired } if !user.CheckPassword(password) { return ErrPasswordIncorrect } } // 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 } // CreateLoginSession creates a temporary login session for 2FA func (s *TotpService) CreateLoginSession(ctx context.Context, userID int64, email string) (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), } 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 func (s *TotpService) GetVerificationMethod(ctx context.Context) *VerificationMethod { if s.settingService.IsEmailVerifyEnabled(ctx) { return &VerificationMethod{Method: "email"} } return &VerificationMethod{Method: "password"} } // SendVerifyCode sends an email verification code for TOTP operations func (s *TotpService) SendVerifyCode(ctx context.Context, userID int64) 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) }