package handlers import ( "fmt" "net/http" "strings" "github.com/gin-gonic/gin" "github.com/pquerna/otp/totp" "github.com/starfleetcptn/gomft/components" "github.com/starfleetcptn/gomft/internal/auth" "golang.org/x/crypto/bcrypt" ) // Handle2FASetup handles the GET /profile/2fa/setup route func (h *Handlers) Handle2FASetup(c *gin.Context) { // Get user from context userID := c.GetUint("userID") var user struct { Email string TwoFactorEnabled bool } if err := h.DB.Table("users").Select("email, two_factor_enabled").Where("id = ?", userID).First(&user).Error; err != nil { c.String(http.StatusInternalServerError, "Failed to get user") return } // Check if 2FA is already enabled if user.TwoFactorEnabled { c.Redirect(http.StatusFound, "/profile") return } // Generate TOTP secret and QR code URL secret, qrCodeURL, err := auth.GenerateTOTPSecret(user.Email) if err != nil { c.String(http.StatusInternalServerError, "Failed to generate 2FA secret") return } // Generate backup codes - now returns both plain codes and hashed codes backupCodesPlain, backupCodesHashed, err := auth.GenerateBackupCodes() if err != nil { c.String(http.StatusInternalServerError, "Failed to generate backup codes") return } // Store secret and backup codes in session temporarily // We store the plain secret in the cookie since it's temporary and will be encrypted before DB storage c.SetCookie("2fa_setup_secret", secret, 3600, "/", "", false, true) c.SetCookie("2fa_setup_backup_codes_hashed", backupCodesHashed, 3600, "/", "", false, true) // Render setup page data := components.TwoFactorSetupData{ QRCodeURL: qrCodeURL, Secret: secret, BackupCodes: backupCodesPlain, // Show plain codes to the user ErrorMessage: "", } components.TwoFactorSetup(c.Request.Context(), data).Render(c, c.Writer) } // Handle2FAVerifySetup handles the POST /profile/2fa/verify route func (h *Handlers) Handle2FAVerifySetup(c *gin.Context) { // Get user from context userID := c.GetUint("userID") var user struct { Email string } if err := h.DB.Table("users").Select("email").Where("id = ?", userID).First(&user).Error; err != nil { c.String(http.StatusInternalServerError, "Failed to get user") return } // Get secret from session secret, err := c.Cookie("2fa_setup_secret") if err != nil { c.String(http.StatusBadRequest, "Setup session expired") return } // Get backup codes from session - now using the hashed version backupCodesHashed, err := c.Cookie("2fa_setup_backup_codes_hashed") if err != nil { c.String(http.StatusBadRequest, "Setup session expired") return } // Verify the code code := c.PostForm("code") // For verification during setup, we use the plain secret since it's not yet encrypted if !totp.Validate(code, secret) { // Regenerate QR code URL using the existing secret qrCodeURL, err := auth.GenerateQRCodeURL(secret, user.Email) if err != nil { c.String(http.StatusInternalServerError, "Failed to generate QR code") return } // For display, we need to generate new plain-text codes // but we'll keep the same hashed codes for storage backupCodesPlain := []string{} if backupCodesHashed != "" { // Create placeholder codes since we can't recover the original codes // We'll use placeholder text that indicates these were already generated for i := 0; i < auth.BackupCodeCount; i++ { backupCodesPlain = append(backupCodesPlain, "[BACKUP CODE ALREADY GENERATED]") } } data := components.TwoFactorSetupData{ QRCodeURL: qrCodeURL, Secret: secret, BackupCodes: backupCodesPlain, ErrorMessage: "Invalid verification code. Please try again.", } components.TwoFactorSetup(c.Request.Context(), data).Render(c, c.Writer) return } // Encrypt the secret before storing in database encryptedSecret, err := auth.EncryptTOTPSecret(secret) if err != nil { c.String(http.StatusInternalServerError, "Failed to secure 2FA secret") return } // Update user with 2FA settings if err := h.DB.Table("users").Where("id = ?", userID).Updates(map[string]interface{}{ "two_factor_secret": encryptedSecret, // Store the encrypted secret "two_factor_enabled": true, "backup_codes": backupCodesHashed, // Store the hashed codes }).Error; err != nil { c.String(http.StatusInternalServerError, "Failed to enable 2FA") return } // Clear setup cookies c.SetCookie("2fa_setup_secret", "", -1, "/", "", false, true) c.SetCookie("2fa_setup_backup_codes_hashed", "", -1, "/", "", false, true) // Redirect to profile with success message c.Redirect(http.StatusFound, "/profile?message=2FA+enabled+successfully") } // Handle2FAVerifyPage handles the GET /login/verify route func (h *Handlers) Handle2FAVerifyPage(c *gin.Context) { // Check if we have a temporary user ID _, err := c.Cookie("temp_user_id") if err != nil { c.Redirect(http.StatusFound, "/login") return } // Render verification page data := components.TwoFactorVerifyData{ ErrorMessage: "", } components.TwoFactorVerify(c.Request.Context(), data).Render(c, c.Writer) } // Handle2FAVerify handles the POST /login/verify route func (h *Handlers) Handle2FAVerify(c *gin.Context) { // Get user ID from cookie tempUserID, err := c.Cookie("temp_user_id") if err != nil { c.Redirect(http.StatusFound, "/login") return } // Parse user ID var userID uint if _, err := fmt.Sscanf(tempUserID, "%d", &userID); err != nil { c.Redirect(http.StatusFound, "/login") return } var user struct { TwoFactorSecret string BackupCodes string Email string IsAdmin *bool } if err := h.DB.Table("users").Select("two_factor_secret, backup_codes, email, is_admin").Where("id = ?", userID).First(&user).Error; err != nil { c.Redirect(http.StatusFound, "/login") return } code := c.PostForm("code") // First try TOTP code if auth.ValidateTOTPCode(user.TwoFactorSecret, code) { // Generate new JWT token and set cookie isAdmin := false if user.IsAdmin != nil { isAdmin = *user.IsAdmin } token, err := h.GenerateJWT(userID, user.Email, isAdmin) if err != nil { c.String(http.StatusInternalServerError, "Failed to generate token") return } c.SetCookie("jwt_token", token, 86400, "/", "", false, true) // Clear temporary user ID cookie c.SetCookie("temp_user_id", "", -1, "/", "", false, true) c.Redirect(http.StatusFound, "/dashboard") return } // Then try backup code if auth.ValidateBackupCode(code, user.BackupCodes) { // Remove used backup code newBackupCodes := auth.RemoveBackupCode(code, user.BackupCodes) if err := h.DB.Table("users").Where("id = ?", userID).Update("backup_codes", newBackupCodes).Error; err != nil { c.String(http.StatusInternalServerError, "Failed to update backup codes") return } // Generate new JWT token and set cookie isAdmin := false if user.IsAdmin != nil { isAdmin = *user.IsAdmin } token, err := h.GenerateJWT(userID, user.Email, isAdmin) if err != nil { c.String(http.StatusInternalServerError, "Failed to generate token") return } c.SetCookie("jwt_token", token, 86400, "/", "", false, true) // Clear temporary user ID cookie c.SetCookie("temp_user_id", "", -1, "/", "", false, true) c.Redirect(http.StatusFound, "/dashboard") return } // If neither code is valid, show error data := components.TwoFactorVerifyData{ ErrorMessage: "Invalid verification code. Please try again.", } components.TwoFactorVerify(c.Request.Context(), data).Render(c, c.Writer) } // Handle2FADisable handles the POST /profile/2fa/disable route func (h *Handlers) Handle2FADisable(c *gin.Context) { // Get user ID from context userID := c.GetUint("userID") // Get current password from form currentPassword := c.PostForm("current_password") if currentPassword == "" { c.Data(http.StatusBadRequest, "text/html", []byte(``)) return } // Get user from database var user struct { PasswordHash string TwoFactorEnabled bool } if err := h.DB.Table("users").Select("password_hash, two_factor_enabled").Where("id = ?", userID).First(&user).Error; err != nil { c.Data(http.StatusInternalServerError, "text/html", []byte(``)) return } // Verify current password if err := bcrypt.CompareHashAndPassword([]byte(user.PasswordHash), []byte(currentPassword)); err != nil { c.Data(http.StatusBadRequest, "text/html", []byte(``)) return } // Check if 2FA is already disabled if !user.TwoFactorEnabled { c.Data(http.StatusBadRequest, "text/html", []byte(``)) return } // Disable 2FA if err := h.DB.Table("users").Where("id = ?", userID).Updates(map[string]interface{}{ "two_factor_enabled": false, "two_factor_secret": nil, "backup_codes": nil, }).Error; err != nil { c.Data(http.StatusInternalServerError, "text/html", []byte(``)) return } // Return success message c.Data(http.StatusOK, "text/html", []byte(``)) } // Handle2FABackupCodes handles the GET /profile/2fa/backup-codes route func (h *Handlers) Handle2FABackupCodes(c *gin.Context) { // Get user from context userID := c.GetUint("userID") var user struct { BackupCodes string TwoFactorEnabled bool Email string // We need email to generate new backup codes display } if err := h.DB.Table("users").Select("backup_codes, two_factor_enabled, email").Where("id = ?", userID).First(&user).Error; err != nil { c.String(http.StatusInternalServerError, "Failed to get user") return } // Check if 2FA is enabled if !user.TwoFactorEnabled { c.Redirect(http.StatusFound, "/profile") return } // If user just generated new codes, we need to show them // We can't derive the plain codes from the hashed ones, so we'll generate new ones for display message := c.Query("message") successMessage := "" if message != "" { successMessage = message } // For backup codes display, we need to check where we are in the flow var backupCodes []string // If the user just regenerated codes or they're viewing for the first time // we need to generate new codes to display, as we can't recover the hashed ones // We'll generate new temporary codes for display only, keeping the hashed ones in the database if strings.Contains(successMessage, "New backup codes generated") { // Generate new codes to display newCodes, _, err := auth.GenerateBackupCodes() if err != nil { c.String(http.StatusInternalServerError, "Failed to generate backup codes for display") return } backupCodes = newCodes // Add a special warning about these being the only time they'll see these codes successMessage = "New backup codes generated successfully. IMPORTANT: Save these codes now. You won't be able to see them again!" } else { // If they're just viewing an existing page, show a message explaining // that backup codes are securely stored and they need to regenerate to view new ones backupCodes = []string{} if user.BackupCodes != "" { // Count how many backup codes the user has by counting commas+1 codeCount := 1 if user.BackupCodes != "" { codeCount = strings.Count(user.BackupCodes, ",") + 1 } // Show a placeholder message for existing codes for i := 0; i < codeCount; i++ { backupCodes = append(backupCodes, "[REDACTED FOR SECURITY]") } // Set a message explaining why codes are hidden successMessage = "For security, backup codes are not displayed after initial generation. Generate new codes to replace existing ones." } } // Render backup codes page data := components.TwoFactorBackupCodesData{ BackupCodes: backupCodes, SuccessMessage: successMessage, } components.TwoFactorBackupCodes(c.Request.Context(), data).Render(c, c.Writer) } // Handle2FARegenerateCodes handles the POST /profile/2fa/regenerate-codes route func (h *Handlers) Handle2FARegenerateCodes(c *gin.Context) { // Get user from context userID := c.GetUint("userID") var user struct { TwoFactorEnabled bool } if err := h.DB.Table("users").Select("two_factor_enabled").Where("id = ?", userID).First(&user).Error; err != nil { c.String(http.StatusInternalServerError, "Failed to get user") return } // Check if 2FA is enabled if !user.TwoFactorEnabled { c.Redirect(http.StatusFound, "/profile") return } // Generate new backup codes - we only need the hashed version for storage _, backupCodesHashed, err := auth.GenerateBackupCodes() if err != nil { c.String(http.StatusInternalServerError, "Failed to generate backup codes") return } // Store new backup codes (hashed version) if err := h.DB.Table("users").Where("id = ?", userID).Update("backup_codes", backupCodesHashed).Error; err != nil { c.String(http.StatusInternalServerError, "Failed to update backup codes") return } // Redirect back to backup codes page with success message c.Redirect(http.StatusFound, "/profile/2fa/backup-codes?message=New backup codes generated successfully") } // Handle2FABackupCodePage handles the GET /login/backup-code route func (h *Handlers) Handle2FABackupCodePage(c *gin.Context) { // Check if we have a temporary user ID _, err := c.Cookie("temp_user_id") if err != nil { c.Redirect(http.StatusFound, "/login") return } // Render backup code verification page data := components.BackupCodeVerifyData{ ErrorMessage: "", } components.TwoFactorBackupVerify(c.Request.Context(), data).Render(c, c.Writer) }