feat: Implement Two-Factor Authentication (2FA) functionality

- Added support for Two-Factor Authentication, including setup, verification, and disabling features.
- Introduced new components for 2FA setup and verification, enhancing user security.
- Updated user model to include fields for 2FA secret, status, and backup codes.
- Created migration to add necessary database fields for 2FA.
- Enhanced authentication handlers to manage 2FA processes during login and profile updates.
- Updated routes to include 2FA-related endpoints for setup and verification.

internal/web/handlers/auth_handlers.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"crypto/rand"
 	"encoding/base64"
+	"fmt"
 	"log"
 	"net/http"
 	"strings"
@@ -56,11 +57,21 @@ func (h *Handlers) AuthMiddleware() gin.HandlerFunc {
 			return
 		}
 
+		// Safely extract claims with type assertions and defaults
+		userID, _ := claims["user_id"].(float64)
+		email, _ := claims["email"].(string)
+		username, _ := claims["username"].(string)
+		isAdmin, _ := claims["is_admin"].(bool)
+
 		// Set user information in the context
-		c.Set("userID", uint(claims["user_id"].(float64)))
-		c.Set("email", claims["email"].(string))
-		c.Set("username", claims["username"].(string))
-		c.Set("isAdmin", claims["is_admin"].(bool))
+		c.Set("userID", uint(userID))
+		if email != "" {
+			c.Set("email", email)
+		}
+		if username != "" {
+			c.Set("username", username)
+		}
+		c.Set("isAdmin", isAdmin)
 
 		c.Next()
 	}
@@ -142,10 +153,11 @@ func (h *Handlers) APIAdminMiddleware() gin.HandlerFunc {
 }
 
 // GenerateJWT generates a JWT token for the given user
-func (h *Handlers) GenerateJWT(userID uint, username string, isAdmin bool) (string, error) {
+func (h *Handlers) GenerateJWT(userID uint, email string, isAdmin bool) (string, error) {
 	token := jwt.NewWithClaims(jwt.SigningMethodHS256, jwt.MapClaims{
 		"user_id":  userID,
-		"username": username,
+		"email":    email,
+		"username": strings.Split(email, "@")[0], // Use email prefix as username
 		"is_admin": isAdmin,
 		"exp":      time.Now().Add(time.Hour * 24).Unix(),
 	})
@@ -243,24 +255,30 @@ func (h *Handlers) HandleLogin(c *gin.Context) {
 		return
 	}
 
-	// Generate JWT token with all necessary user information
-	token := jwt.NewWithClaims(jwt.SigningMethodHS256, jwt.MapClaims{
-		"user_id":  user.ID,
-		"email":    user.Email,
-		"username": strings.Split(user.Email, "@")[0], // Use email prefix as username
-		"is_admin": user.IsAdmin,
-		"exp":      time.Now().Add(time.Hour * 24).Unix(),
-	})
+	// Check if 2FA is enabled
+	if user.TwoFactorEnabled {
+		// Store user ID temporarily for 2FA verification
+		c.SetCookie("temp_user_id", fmt.Sprintf("%d", user.ID), 300, "/", "", false, true) // 5 minutes expiry
 
-	// Sign the token
-	tokenString, err := token.SignedString([]byte(h.JWTSecret))
+		// Redirect to 2FA verification page
+		c.Redirect(http.StatusFound, "/login/verify")
+		return
+	}
+
+	// If 2FA is not enabled, proceed with normal login
+	// Generate JWT token with all necessary user information
+	isAdmin := false
+	if user.IsAdmin != nil {
+		isAdmin = *user.IsAdmin
+	}
+	token, err := h.GenerateJWT(user.ID, user.Email, isAdmin)
 	if err != nil {
 		components.Login(components.CreateTemplateContext(c), "Authentication error").Render(c, c.Writer)
 		return
 	}
 
 	// Set token in cookie
-	c.SetCookie("jwt_token", tokenString, 86400, "/", "", false, true)
+	c.SetCookie("jwt_token", token, 86400, "/", "", false, true)
 	c.Redirect(http.StatusFound, "/dashboard")
 }
 

internal/web/handlers/routes.go

@@ -10,6 +10,8 @@ func (h *Handlers) RegisterRoutes(router *gin.Engine) {
 	router.GET("/", h.HandleHome)
 	router.GET("/login", h.HandleLoginPage)
 	router.POST("/login", h.HandleLogin)
+	router.GET("/login/verify", h.Handle2FAVerifyPage)
+	router.POST("/login/verify", h.Handle2FAVerify)
 	router.GET("/forgot-password", h.HandleForgotPasswordPage)
 	router.POST("/forgot-password", h.HandleForgotPassword)
 	router.GET("/reset-password", h.HandleResetPasswordPage)
@@ -22,6 +24,11 @@ func (h *Handlers) RegisterRoutes(router *gin.Engine) {
 	// Password change route - only accessed from profile page
 	authorized.POST("/change-password", h.HandleChangePassword)
 
+	// 2FA routes - under profile
+	authorized.GET("/profile/2fa/setup", h.Handle2FASetup)
+	authorized.POST("/profile/2fa/verify", h.Handle2FAVerifySetup)
+	authorized.POST("/profile/2fa/disable", h.Handle2FADisable)
+
 	{
 		authorized.GET("/dashboard", h.HandleDashboard)
 		authorized.GET("/configs", h.HandleConfigs)

internal/web/handlers/two_factor_handlers.go

@@ -0,0 +1,292 @@
+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(`<div class="bg-red-100 border border-red-400 text-red-700 px-4 py-3 rounded mb-4" role="alert">
+			<span class="block sm:inline">Current password is required</span>
+		</div>`))
+		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(`<div class="bg-red-100 border border-red-400 text-red-700 px-4 py-3 rounded mb-4" role="alert">
+			<span class="block sm:inline">Failed to get user information</span>
+		</div>`))
+		return
+	}
+
+	// Verify current password
+	if err := bcrypt.CompareHashAndPassword([]byte(user.PasswordHash), []byte(currentPassword)); err != nil {
+		c.Data(http.StatusBadRequest, "text/html", []byte(`<div class="bg-red-100 border border-red-400 text-red-700 px-4 py-3 rounded mb-4" role="alert">
+			<span class="block sm:inline">Current password is incorrect</span>
+		</div>`))
+		return
+	}
+
+	// Check if 2FA is already disabled
+	if !user.TwoFactorEnabled {
+		c.Data(http.StatusBadRequest, "text/html", []byte(`<div class="bg-yellow-100 border border-yellow-400 text-yellow-700 px-4 py-3 rounded mb-4" role="alert">
+			<span class="block sm:inline">Two-factor authentication is already disabled</span>
+		</div>`))
+		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(`<div class="bg-red-100 border border-red-400 text-red-700 px-4 py-3 rounded mb-4" role="alert">
+			<span class="block sm:inline">Failed to disable two-factor authentication</span>
+		</div>`))
+		return
+	}
+
+	// Return success message
+	c.Data(http.StatusOK, "text/html", []byte(`<div class="bg-green-100 border border-green-400 text-green-700 px-4 py-3 rounded mb-4" role="alert">
+		<span class="block sm:inline">Two-factor authentication has been disabled</span>
+		<script>
+			setTimeout(function() {
+				window.location.reload();
+			}, 1500);
+		</script>
+	</div>`))
+}