package user import ( "context" "crypto/rand" "fmt" "strings" "time" "github.com/pquerna/otp" "github.com/pquerna/otp/totp" "github.com/yaoapp/gou/model" "github.com/yaoapp/kun/maps" "github.com/yaoapp/yao/openapi/oauth/types" "golang.org/x/crypto/bcrypt" ) // User MFA Management // GenerateMFASecret generates a new TOTP secret for user func (u *DefaultUser) GenerateMFASecret(ctx context.Context, userID string, options *types.MFAOptions) (string, string, error) { // Verify user exists m := model.Select(u.model) users, err := m.Get(model.QueryParam{ Select: []interface{}{"user_id", "mfa_enabled"}, Wheres: []model.QueryWhere{ {Column: "user_id", Value: userID}, }, Limit: 1, }) if err != nil { return "", "", fmt.Errorf(ErrFailedToGetUser, err) } if len(users) == 0 { return "", "", fmt.Errorf(ErrUserNotFound) } // Use provided options or fallback to instance defaults if options == nil { options = u.mfaOptions } // Apply defaults for individual fields if not specified issuer := options.Issuer if issuer == "" { issuer = u.mfaOptions.Issuer } accountName := options.AccountName if accountName == "" { accountName = userID // Default to userID } algorithm := options.Algorithm if algorithm == "" { algorithm = u.mfaOptions.Algorithm } digits := options.Digits if digits == 0 { digits = u.mfaOptions.Digits } period := options.Period if period == 0 { period = u.mfaOptions.Period } secretSize := options.SecretSize if secretSize == 0 { secretSize = u.mfaOptions.SecretSize } // Convert algorithm string to otp.Algorithm var otpAlgorithm otp.Algorithm switch algorithm { case "SHA1": otpAlgorithm = otp.AlgorithmSHA1 case "SHA256": otpAlgorithm = otp.AlgorithmSHA256 case "SHA512": otpAlgorithm = otp.AlgorithmSHA512 default: otpAlgorithm = otp.AlgorithmSHA256 // Default fallback } // Convert digits to otp.Digits var otpDigits otp.Digits if digits == 8 { otpDigits = otp.DigitsEight } else { otpDigits = otp.DigitsSix // Default } // Generate TOTP key key, err := totp.Generate(totp.GenerateOpts{ Issuer: issuer, AccountName: accountName, SecretSize: uint(secretSize), Algorithm: otpAlgorithm, Digits: otpDigits, Period: uint(period), }) if err != nil { return "", "", fmt.Errorf(ErrFailedToGenerateMFASecret, err) } secret := key.Secret() qrCodeURL := key.URL() // Store the secret temporarily (not enabled yet until user verifies) updateData := maps.MapStrAny{ "mfa_secret": secret, "mfa_issuer": issuer, "mfa_algorithm": algorithm, "mfa_digits": digits, "mfa_period": period, } _, err = m.UpdateWhere(model.QueryParam{ Wheres: []model.QueryWhere{ {Column: "user_id", Value: userID}, }, Limit: 1, }, updateData) if err != nil { return "", "", fmt.Errorf(ErrFailedToUpdateMFAStatus, err) } return secret, qrCodeURL, nil } // EnableMFA enables multi-factor authentication for user func (u *DefaultUser) EnableMFA(ctx context.Context, userID string, secret string, code string) error { // Get user and current MFA status m := model.Select(u.model) users, err := m.Get(model.QueryParam{ Select: []interface{}{"user_id", "mfa_enabled", "mfa_secret"}, Wheres: []model.QueryWhere{ {Column: "user_id", Value: userID}, }, Limit: 1, }) if err != nil { return fmt.Errorf(ErrFailedToGetUser, err) } if len(users) == 0 { return fmt.Errorf(ErrUserNotFound) } user := users[0] // Check if MFA is already enabled if mfaEnabled, ok := user["mfa_enabled"].(bool); ok && mfaEnabled { return fmt.Errorf(ErrMFAAlreadyEnabled) } // Handle different boolean types from database if mfaEnabledInt, ok := user["mfa_enabled"].(int64); ok && mfaEnabledInt != 0 { return fmt.Errorf(ErrMFAAlreadyEnabled) } // Use stored secret if not provided if secret == "" { if storedSecret, ok := user["mfa_secret"].(string); ok && storedSecret != "" { secret = storedSecret } else { return fmt.Errorf("no MFA secret found, please generate one first") } } // Verify the provided code valid := totp.Validate(code, secret) if !valid { return fmt.Errorf(ErrInvalidMFACode) } // Enable MFA updateData := maps.MapStrAny{ "mfa_enabled": true, "mfa_secret": secret, // Store the verified secret "mfa_enabled_at": time.Now(), } affected, err := m.UpdateWhere(model.QueryParam{ Wheres: []model.QueryWhere{ {Column: "user_id", Value: userID}, }, Limit: 1, }, updateData) if err != nil { return fmt.Errorf(ErrFailedToUpdateMFAStatus, err) } if affected == 0 { // Check if user exists exists, checkErr := u.UserExists(ctx, userID) if checkErr != nil { return fmt.Errorf(ErrFailedToUpdateMFAStatus, checkErr) } if !exists { return fmt.Errorf(ErrUserNotFound) } // User exists but no changes were made (already enabled with same secret) } return nil } // DisableMFA disables multi-factor authentication for user func (u *DefaultUser) DisableMFA(ctx context.Context, userID string, code string) error { // Get user and current MFA status m := model.Select(u.model) users, err := m.Get(model.QueryParam{ Select: []interface{}{"user_id", "mfa_enabled", "mfa_secret"}, Wheres: []model.QueryWhere{ {Column: "user_id", Value: userID}, }, Limit: 1, }) if err != nil { return fmt.Errorf(ErrFailedToGetUser, err) } if len(users) == 0 { return fmt.Errorf(ErrUserNotFound) } user := users[0] // Check if MFA is enabled mfaEnabled := false if enabled, ok := user["mfa_enabled"].(bool); ok { mfaEnabled = enabled } else if enabledInt, ok := user["mfa_enabled"].(int64); ok { mfaEnabled = enabledInt != 0 } if !mfaEnabled { return fmt.Errorf(ErrMFANotEnabled) } // Get stored secret secret, ok := user["mfa_secret"].(string) if !ok || secret == "" { return fmt.Errorf("no MFA secret found") } // Verify the provided code valid := totp.Validate(code, secret) if !valid { return fmt.Errorf(ErrInvalidMFACode) } // Disable MFA and clear sensitive data updateData := maps.MapStrAny{ "mfa_enabled": false, "mfa_secret": nil, // Clear the secret "mfa_recovery_hash": nil, // Clear recovery codes "mfa_enabled_at": nil, "mfa_last_verified_at": nil, } affected, err := m.UpdateWhere(model.QueryParam{ Wheres: []model.QueryWhere{ {Column: "user_id", Value: userID}, }, Limit: 1, }, updateData) if err != nil { return fmt.Errorf(ErrFailedToUpdateMFAStatus, err) } if affected == 0 { // Check if user exists exists, checkErr := u.UserExists(ctx, userID) if checkErr != nil { return fmt.Errorf(ErrFailedToUpdateMFAStatus, checkErr) } if !exists { return fmt.Errorf(ErrUserNotFound) } // User exists but no changes were made (already disabled) } return nil } // VerifyMFACode verifies a TOTP code for user func (u *DefaultUser) VerifyMFACode(ctx context.Context, userID string, code string) (bool, error) { // Get user and MFA status m := model.Select(u.model) users, err := m.Get(model.QueryParam{ Select: []interface{}{"user_id", "mfa_enabled", "mfa_secret"}, Wheres: []model.QueryWhere{ {Column: "user_id", Value: userID}, }, Limit: 1, }) if err != nil { return false, fmt.Errorf(ErrFailedToGetUser, err) } if len(users) == 0 { return false, fmt.Errorf(ErrUserNotFound) } user := users[0] // Check if MFA is enabled mfaEnabled := false if enabled, ok := user["mfa_enabled"].(bool); ok { mfaEnabled = enabled } else if enabledInt, ok := user["mfa_enabled"].(int64); ok { mfaEnabled = enabledInt != 0 } if !mfaEnabled { return false, fmt.Errorf(ErrMFANotEnabled) } // Get stored secret secret, ok := user["mfa_secret"].(string) if !ok || secret == "" { return false, fmt.Errorf("no MFA secret found") } // Verify the code valid := totp.Validate(code, secret) if valid { // Update last verified timestamp _, err = m.UpdateWhere(model.QueryParam{ Wheres: []model.QueryWhere{ {Column: "user_id", Value: userID}, }, Limit: 1, }, maps.MapStrAny{ "mfa_last_verified_at": time.Now(), }) // Don't fail verification if timestamp update fails if err != nil { // Log the error but continue } } return valid, nil } // GenerateRecoveryCodes generates new recovery codes for user and stores their hash func (u *DefaultUser) GenerateRecoveryCodes(ctx context.Context, userID string) ([]string, error) { // Verify user exists and MFA is enabled m := model.Select(u.model) users, err := m.Get(model.QueryParam{ Select: []interface{}{"user_id", "mfa_enabled"}, Wheres: []model.QueryWhere{ {Column: "user_id", Value: userID}, }, Limit: 1, }) if err != nil { return nil, fmt.Errorf(ErrFailedToGetUser, err) } if len(users) == 0 { return nil, fmt.Errorf(ErrUserNotFound) } user := users[0] // Check if MFA is enabled mfaEnabled := false if enabled, ok := user["mfa_enabled"].(bool); ok { mfaEnabled = enabled } else if enabledInt, ok := user["mfa_enabled"].(int64); ok { mfaEnabled = enabledInt != 0 } if !mfaEnabled { return nil, fmt.Errorf(ErrMFANotEnabled) } // Generate multiple recovery codes and store bcrypt hashes (512 char limit) recoveryCount := u.mfaOptions.RecoveryCount recoveryLength := u.mfaOptions.RecoveryLength recoveryCodes := make([]string, recoveryCount) for i := 0; i < recoveryCount; i++ { code, err := generateRecoveryCode(recoveryLength) if err != nil { return nil, fmt.Errorf("failed to generate recovery code: %w", err) } recoveryCodes[i] = code } // Hash each recovery code with bcrypt and store all hashes recoveryHashes := make([]string, recoveryCount) for i, code := range recoveryCodes { hashedCode, err := bcrypt.GenerateFromPassword([]byte(code), bcrypt.DefaultCost) if err != nil { return nil, fmt.Errorf("failed to hash recovery code: %w", err) } recoveryHashes[i] = string(hashedCode) } // Join all bcrypt hashes (~60 bytes each, 8 hashes = ~480 bytes, under 512 limit) allHashesStr := strings.Join(recoveryHashes, "|||") updateData := maps.MapStrAny{ "mfa_recovery_hash": allHashesStr, // Store bcrypt hashes (~480 bytes, under 512 limit) } affected, err := m.UpdateWhere(model.QueryParam{ Wheres: []model.QueryWhere{ {Column: "user_id", Value: userID}, }, Limit: 1, }, updateData) if err != nil { return nil, fmt.Errorf(ErrFailedToUpdateMFAStatus, err) } if affected == 0 { // Check if user exists exists, checkErr := u.UserExists(ctx, userID) if checkErr != nil { return nil, fmt.Errorf(ErrFailedToUpdateMFAStatus, checkErr) } if !exists { return nil, fmt.Errorf(ErrUserNotFound) } // User exists but no changes were made } // Return all generated recovery codes return recoveryCodes, nil } // VerifyRecoveryCode verifies and consumes a recovery code func (u *DefaultUser) VerifyRecoveryCode(ctx context.Context, userID string, code string) (bool, error) { // Get user and MFA status m := model.Select(u.model) users, err := m.Get(model.QueryParam{ Select: []interface{}{"user_id", "mfa_enabled", "mfa_recovery_hash"}, Wheres: []model.QueryWhere{ {Column: "user_id", Value: userID}, }, Limit: 1, }) if err != nil { return false, fmt.Errorf(ErrFailedToGetUser, err) } if len(users) == 0 { return false, fmt.Errorf(ErrUserNotFound) } user := users[0] // Check if MFA is enabled mfaEnabled := false if enabled, ok := user["mfa_enabled"].(bool); ok { mfaEnabled = enabled } else if enabledInt, ok := user["mfa_enabled"].(int64); ok { mfaEnabled = enabledInt != 0 } if !mfaEnabled { return false, fmt.Errorf(ErrMFANotEnabled) } // Get stored bcrypt hashes string (512 char limit - no problem!) recoveryHashesStr, ok := user["mfa_recovery_hash"].(string) if !ok || recoveryHashesStr == "" { return false, fmt.Errorf("no recovery codes found") } // Split into hashes list recoveryHashes := strings.Split(recoveryHashesStr, "|||") if len(recoveryHashes) == 0 { return false, fmt.Errorf("no recovery codes found") } // Check if user input code matches any stored bcrypt hash matchIndex := -1 for i, storedHash := range recoveryHashes { if storedHash == "" { continue // Skip already used codes } // Verify user input against stored bcrypt hash err := bcrypt.CompareHashAndPassword([]byte(storedHash), []byte(code)) if err == nil { matchIndex = i break } } if matchIndex == -1 { return false, nil // Invalid code } // Mark the code as used by clearing its hash recoveryHashes[matchIndex] = "" updatedHashesStr := strings.Join(recoveryHashes, "|||") // Update recovery codes in database and mark verification time updateData := maps.MapStrAny{ "mfa_recovery_hash": updatedHashesStr, "mfa_last_verified_at": time.Now(), } _, err = m.UpdateWhere(model.QueryParam{ Wheres: []model.QueryWhere{ {Column: "user_id", Value: userID}, }, Limit: 1, }, updateData) if err != nil { return false, fmt.Errorf(ErrFailedToUpdateMFAStatus, err) } return true, nil } // IsMFAEnabled checks if MFA is enabled for a user func (u *DefaultUser) IsMFAEnabled(ctx context.Context, userID string) (bool, error) { m := model.Select(u.model) users, err := m.Get(model.QueryParam{ Select: []interface{}{"user_id", "mfa_enabled"}, Wheres: []model.QueryWhere{ {Column: "user_id", Value: userID}, }, Limit: 1, }) if err != nil { return false, fmt.Errorf(ErrFailedToGetUser, err) } if len(users) == 0 { return false, fmt.Errorf(ErrUserNotFound) } user := users[0] // Check MFA status if enabled, ok := user["mfa_enabled"].(bool); ok { return enabled, nil } // Handle different boolean types from database if enabledInt, ok := user["mfa_enabled"].(int64); ok { return enabledInt != 0, nil } return false, nil } // GetMFAConfig retrieves MFA configuration for a user func (u *DefaultUser) GetMFAConfig(ctx context.Context, userID string) (maps.MapStrAny, error) { m := model.Select(u.model) users, err := m.Get(model.QueryParam{ Select: []interface{}{ "user_id", "mfa_enabled", "mfa_issuer", "mfa_algorithm", "mfa_digits", "mfa_period", "mfa_enabled_at", "mfa_last_verified_at", "mfa_recovery_hash", // Include recovery hash field }, Wheres: []model.QueryWhere{ {Column: "user_id", Value: userID}, }, Limit: 1, }) if err != nil { return nil, fmt.Errorf(ErrFailedToGetUser, err) } if len(users) == 0 { return nil, fmt.Errorf(ErrUserNotFound) } user := users[0] // Check if MFA is enabled mfaEnabled := false if enabled, ok := user["mfa_enabled"].(bool); ok { mfaEnabled = enabled } else if enabledInt, ok := user["mfa_enabled"].(int64); ok { mfaEnabled = enabledInt != 0 } config := maps.MapStrAny{ "user_id": userID, "mfa_enabled": mfaEnabled, } if mfaEnabled { // Include MFA configuration details (but not the secret) config["mfa_issuer"] = user["mfa_issuer"] config["mfa_algorithm"] = user["mfa_algorithm"] config["mfa_digits"] = user["mfa_digits"] config["mfa_period"] = user["mfa_period"] config["mfa_enabled_at"] = user["mfa_enabled_at"] config["mfa_last_verified_at"] = user["mfa_last_verified_at"] // Check how many recovery codes are available (bcrypt hash storage) if recoveryHashesStr, ok := user["mfa_recovery_hash"].(string); ok && recoveryHashesStr != "" { hashes := strings.Split(recoveryHashesStr, "|||") remainingCodes := 0 for _, hash := range hashes { if hash != "" { remainingCodes++ } } config["recovery_codes_available"] = remainingCodes } else { config["recovery_codes_available"] = 0 } } return config, nil } // Helper function to generate recovery codes func generateRecoveryCode(length int) (string, error) { // Use alphanumeric charset (excluding similar-looking characters for better UX) // Excludes: 0, O, 1, I, l to avoid confusion const charset = "23456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnpqrstuvwxyz" b := make([]byte, length) _, err := rand.Read(b) if err != nil { return "", err } for i := range b { b[i] = charset[b[i]%byte(len(charset))] } // Format with dashes for better readability (like GitHub) // For 8-character codes: XXXX-XXXX // For 12-character codes: XXXX-XXXX-XXXX result := string(b) if length == 8 { return fmt.Sprintf("%s-%s", result[:4], result[4:]), nil } else if length >= 12 { return fmt.Sprintf("%s-%s-%s", result[:4], result[4:8], result[8:]), nil } return result, nil }