package handlers import ( "fmt" "net/http" "strings" "github.com/gin-gonic/gin" "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 backupCodes, err := auth.GenerateBackupCodes() if err != nil { c.String(http.StatusInternalServerError, "Failed to generate backup codes") return } // Store secret and backup codes in session temporarily c.SetCookie("2fa_setup_secret", secret, 3600, "/", "", false, true) c.SetCookie("2fa_setup_backup_codes", strings.Join(backupCodes, ","), 3600, "/", "", false, true) // Render setup page data := components.TwoFactorSetupData{ QRCodeURL: qrCodeURL, Secret: secret, BackupCodes: backupCodes, 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 backupCodes, err := c.Cookie("2fa_setup_backup_codes") if err != nil { c.String(http.StatusBadRequest, "Setup session expired") return } // Verify the code code := c.PostForm("code") if !auth.ValidateTOTPCode(secret, code) { // 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 } data := components.TwoFactorSetupData{ QRCodeURL: qrCodeURL, Secret: secret, BackupCodes: strings.Split(backupCodes, ","), ErrorMessage: "Invalid verification code. Please try again.", } components.TwoFactorSetup(c.Request.Context(), data).Render(c, c.Writer) return } // Update user with 2FA settings if err := h.DB.Table("users").Where("id = ?", userID).Updates(map[string]interface{}{ "two_factor_secret": secret, "two_factor_enabled": true, "backup_codes": backupCodes, }).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", "", -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.Model("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(``)) }