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(`