diff --git a/components/profile.templ b/components/profile.templ index 25b9d03..700fa79 100644 --- a/components/profile.templ +++ b/components/profile.templ @@ -5,8 +5,112 @@ import ( "github.com/starfleetcptn/gomft/internal/db" ) +// Dialog component for 2FA disable confirmation +templ TwoFactorDisableDialog() { + +} + templ Profile(ctx context.Context, user db.User) { @LayoutWithContext("Profile", ctx) { + + + @TwoFactorDisableDialog()

@@ -45,6 +149,35 @@ templ Profile(ctx context.Context, user db.User) { }

+
+
Two-Factor Authentication
+
+ if user.TwoFactorEnabled { +
+ + Enabled + + +
+ } else { +
+ + Disabled + + + + Enable 2FA + +
+ } +
+
Theme
diff --git a/components/providers/destination/gphotos.templ b/components/providers/destination/gphotos.templ index 4398c37..637c99d 100644 --- a/components/providers/destination/gphotos.templ +++ b/components/providers/destination/gphotos.templ @@ -74,6 +74,7 @@ templ GooglePhotosDestinationForm() {
+
+
+
+

Set Up Two-Factor Authentication

+

Enhance your account security with 2FA

+
+ + if data.ErrorMessage != "" { + + } + +
+
+

1. Scan QR Code

+

+ Scan this QR code with your authenticator app (Google Authenticator, Authy, etc.) +

+
+ QR Code +
+
+

+ Can't scan the QR code? Use this code instead: +

+ + { data.Secret } + +
+
+ +
+

2. Verify Setup

+
+
+ + +
+ +
+
+ + if len(data.BackupCodes) > 0 { +
+

3. Save Backup Codes

+

+ Store these backup codes in a safe place. You can use them to access your account if you lose your authenticator device. +

+
+ for _, code := range data.BackupCodes { +
+ { code } +
+ } +
+
+ + +
+
+ } +
+
+
+
+ } +} \ No newline at end of file diff --git a/components/two_factor_verify.templ b/components/two_factor_verify.templ new file mode 100644 index 0000000..08fa043 --- /dev/null +++ b/components/two_factor_verify.templ @@ -0,0 +1,83 @@ +package components + +import "context" + +type TwoFactorVerifyData struct { + ErrorMessage string +} + +templ TwoFactorVerify(ctx context.Context, data TwoFactorVerifyData) { + @LayoutWithContext("Two-Factor Authentication", ctx) { +
+
+
+
+
+ +
+

Two-Factor Authentication

+

Enter the code from your authenticator app

+
+ + if data.ErrorMessage != "" { + + } + +
+
+ +
+
+ +
+ +
+
+ + + +
+ + Use a backup code instead + +
+
+
+
+
+ } +} \ No newline at end of file diff --git a/go.mod b/go.mod index 66e8af5..80fee0b 100644 --- a/go.mod +++ b/go.mod @@ -4,11 +4,13 @@ go 1.24.0 require ( github.com/a-h/templ v0.3.833 + github.com/gin-contrib/sessions v1.0.2 github.com/gin-gonic/gin v1.10.0 github.com/glebarez/sqlite v1.11.0 github.com/go-gormigrate/gormigrate/v2 v2.1.3 github.com/golang-jwt/jwt/v5 v5.2.1 github.com/joho/godotenv v1.5.1 + github.com/pquerna/otp v1.4.0 github.com/robfig/cron/v3 v3.0.1 github.com/stretchr/testify v1.10.0 golang.org/x/crypto v0.35.0 @@ -17,13 +19,13 @@ require ( ) require ( + github.com/boombuler/barcode v1.0.1-0.20190219062509-6c824513bacc // indirect github.com/bytedance/sonic v1.12.9 // indirect github.com/bytedance/sonic/loader v0.2.3 // indirect github.com/cloudwego/base64x v0.1.5 // indirect github.com/davecgh/go-spew v1.1.1 // indirect github.com/dustin/go-humanize v1.0.1 // indirect github.com/gabriel-vasile/mimetype v1.4.8 // indirect - github.com/gin-contrib/sessions v1.0.2 // indirect github.com/gin-contrib/sse v1.0.0 // indirect github.com/glebarez/go-sqlite v1.21.2 // indirect github.com/go-playground/locales v0.14.1 // indirect diff --git a/go.sum b/go.sum index f9cec7f..7d2c387 100644 --- a/go.sum +++ b/go.sum @@ -1,5 +1,7 @@ github.com/a-h/templ v0.3.833 h1:L/KOk/0VvVTBegtE0fp2RJQiBm7/52Zxv5fqlEHiQUU= github.com/a-h/templ v0.3.833/go.mod h1:cAu4AiZhtJfBjMY0HASlyzvkrtjnHWPeEsyGK2YYmfk= +github.com/boombuler/barcode v1.0.1-0.20190219062509-6c824513bacc h1:biVzkmvwrH8WK8raXaxBx6fRVTlJILwEwQGL1I/ByEI= +github.com/boombuler/barcode v1.0.1-0.20190219062509-6c824513bacc/go.mod h1:paBWMcWSl3LHKBqUq+rly7CNSldXjb2rDl3JlRe0mD8= github.com/bytedance/sonic v1.12.9 h1:Od1BvK55NnewtGaJsTDeAOSnLVO2BTSLOe0+ooKokmQ= github.com/bytedance/sonic v1.12.9/go.mod h1:uVvFidNmlt9+wa31S1urfwwthTWteBgG0hWuoKAXTx8= github.com/bytedance/sonic/loader v0.1.1/go.mod h1:ncP89zfokxS5LZrJxl5z0UJcsk4M4yY2JpfqGeCtNLU= @@ -42,6 +44,8 @@ github.com/golang-jwt/jwt/v5 v5.2.1/go.mod h1:pqrtFR0X4osieyHYxtmOUWsAWrfe1Q5UVI github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI= github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg= +github.com/google/gofuzz v1.2.0 h1:xRy4A+RhZaiKjJ1bPfwQ8sedCA+YS2YcCHW6ec7JMi0= +github.com/google/gofuzz v1.2.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg= github.com/google/pprof v0.0.0-20221118152302-e6195bd50e26 h1:Xim43kblpZXfIBQsbuBVKCudVG457BR2GZFIz3uw3hQ= github.com/google/pprof v0.0.0-20221118152302-e6195bd50e26/go.mod h1:dDKJzRmX4S37WGHujM7tX//fmj1uioxKzKxz3lo4HJo= github.com/google/uuid v1.3.0 h1:t6JiXgmwXMjEs8VusXIJk2BXHsn+wx8BZdTaoZ5fu7I= @@ -77,6 +81,8 @@ github.com/pelletier/go-toml/v2 v2.2.3 h1:YmeHyLY8mFWbdkNWwpr+qIL2bEqT0o95WSdkNH github.com/pelletier/go-toml/v2 v2.2.3/go.mod h1:MfCQTFTvCcUyyvvwm1+G6H/jORL20Xlb6rzQu9GuUkc= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/pquerna/otp v1.4.0 h1:wZvl1TIVxKRThZIBiwOOHOGP/1+nZyWBil9Y2XNEDzg= +github.com/pquerna/otp v1.4.0/go.mod h1:dkJfzwRKNiegxyNb54X/3fLwhCynbMspSyWKnvi1AEg= github.com/remyoudompheng/bigfft v0.0.0-20200410134404-eec4a21b6bb0/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo= github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE= github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo= diff --git a/internal/auth/totp.go b/internal/auth/totp.go new file mode 100644 index 0000000..fc46f56 --- /dev/null +++ b/internal/auth/totp.go @@ -0,0 +1,158 @@ +package auth + +import ( + "bytes" + "crypto/rand" + "encoding/base32" + "encoding/base64" + "fmt" + "image/png" + "strings" + + // "github.com/pquerna/otp/base32" + "github.com/pquerna/otp/totp" +) + +const ( + // IssuerName is the name of the issuer that appears in authenticator apps + IssuerName = "GoMFT" + // SecretSize is the size of the TOTP secret in bytes + SecretSize = 20 + // BackupCodeCount is the number of backup codes to generate + BackupCodeCount = 8 + // BackupCodeLength is the length of each backup code + BackupCodeLength = 8 +) + +// GenerateTOTPSecret generates a new TOTP secret for a user +func GenerateTOTPSecret(email string) (string, string, error) { + // Generate TOTP key using the library + key, err := totp.Generate(totp.GenerateOpts{ + Issuer: IssuerName, + AccountName: email, + }) + if err != nil { + return "", "", fmt.Errorf("failed to generate TOTP key: %v", err) + } + + // Generate QR code image + var buf bytes.Buffer + img, err := key.Image(256, 256) + if err != nil { + return "", "", fmt.Errorf("failed to generate QR code image: %v", err) + } + + // Encode image as PNG and convert to base64 + err = png.Encode(&buf, img) + if err != nil { + return "", "", fmt.Errorf("failed to encode QR code image: %v", err) + } + + // Create data URL + dataURL := fmt.Sprintf("data:image/png;base64,%s", base64.StdEncoding.EncodeToString(buf.Bytes())) + + return key.Secret(), dataURL, nil +} + +// ValidateTOTPCode validates a TOTP code against a secret +func ValidateTOTPCode(secret string, code string) bool { + // Remove any spaces from the code + code = strings.ReplaceAll(code, " ", "") + + // Use the library's Validate function + return totp.Validate(code, secret) +} + +// GenerateBackupCodes generates a set of backup codes +func GenerateBackupCodes() ([]string, error) { + codes := make([]string, BackupCodeCount) + for i := 0; i < BackupCodeCount; i++ { + // Generate random bytes + bytes := make([]byte, BackupCodeLength/2) + _, err := rand.Read(bytes) + if err != nil { + return nil, fmt.Errorf("failed to generate backup code: %v", err) + } + + // Convert to hex string + codes[i] = fmt.Sprintf("%x", bytes) + } + return codes, nil +} + +// ValidateBackupCode validates a backup code against a list of codes +func ValidateBackupCode(providedCode string, storedCodes string) bool { + if storedCodes == "" { + return false + } + + // Remove any spaces and convert to lowercase + providedCode = strings.ToLower(strings.ReplaceAll(providedCode, " ", "")) + + // Split stored codes + codes := strings.Split(storedCodes, ",") + + // Check if the provided code matches any stored code + for _, code := range codes { + if code == providedCode { + return true + } + } + + return false +} + +// RemoveBackupCode removes a used backup code from the list +func RemoveBackupCode(usedCode string, storedCodes string) string { + if storedCodes == "" { + return "" + } + + usedCode = strings.ToLower(strings.ReplaceAll(usedCode, " ", "")) + codes := strings.Split(storedCodes, ",") + + var newCodes []string + for _, code := range codes { + if code != usedCode { + newCodes = append(newCodes, code) + } + } + + return strings.Join(newCodes, ",") +} + +// GenerateQRCodeURL generates a QR code URL for an existing secret +func GenerateQRCodeURL(secret string, email string) (string, error) { + // Decode the base32 secret + secretBytes, err := base32.StdEncoding.DecodeString(secret) + if err != nil { + return "", fmt.Errorf("failed to decode secret: %v", err) + } + + key, err := totp.Generate(totp.GenerateOpts{ + Issuer: IssuerName, + AccountName: email, + Secret: secretBytes, + }) + if err != nil { + return "", fmt.Errorf("failed to generate TOTP key: %v", err) + } + + // Generate QR code image + var buf bytes.Buffer + img, err := key.Image(256, 256) + if err != nil { + return "", fmt.Errorf("failed to generate QR code image: %v", err) + } + + // Encode image as PNG and convert to base64 + err = png.Encode(&buf, img) + if err != nil { + return "", fmt.Errorf("failed to encode QR code image: %v", err) + } + + // Create data URL + dataURL := fmt.Sprintf("data:image/png;base64,%s", base64.StdEncoding.EncodeToString(buf.Bytes())) + + return dataURL, nil +} diff --git a/internal/db/db.go b/internal/db/db.go index 0b773bf..a8d9397 100644 --- a/internal/db/db.go +++ b/internal/db/db.go @@ -25,6 +25,9 @@ type User struct { AccountLocked *bool `gorm:"default:false"` LockoutUntil *time.Time Theme string `gorm:"default:'light'"` + TwoFactorSecret string `gorm:"type:varchar(32)"` + TwoFactorEnabled bool `gorm:"default:false"` + BackupCodes string `gorm:"type:text"` // Comma-separated backup codes CreatedAt time.Time UpdatedAt time.Time } diff --git a/internal/db/migrations/003_add_2fa.go b/internal/db/migrations/003_add_2fa.go new file mode 100644 index 0000000..cadc56d --- /dev/null +++ b/internal/db/migrations/003_add_2fa.go @@ -0,0 +1,80 @@ +package migrations + +import ( + "fmt" + "os" + "time" + + "github.com/go-gormigrate/gormigrate/v2" + "gorm.io/gorm" +) + +// Add2FA creates a migration for adding Two-Factor Authentication fields +func Add2FA() *gormigrate.Migration { + return &gormigrate.Migration{ + ID: "003_add_2fa", + Migrate: func(tx *gorm.DB) error { + // Check if any tables exist (indicating an existing database) + var count int64 + if err := tx.Raw("SELECT count(*) FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'").Scan(&count).Error; err != nil { + return fmt.Errorf("failed to check for existing tables: %v", err) + } + + // If tables exist, create a backup + if count > 0 { + // Get the database path + sqlDB, err := tx.DB() + if err != nil { + return fmt.Errorf("failed to get underlying database: %v", err) + } + + var seq int + var name, dbPath string + if err := sqlDB.QueryRow("PRAGMA database_list").Scan(&seq, &name, &dbPath); err != nil { + return fmt.Errorf("failed to get database path: %v", err) + } + + // Create backup file with timestamp + backupFile := fmt.Sprintf("%s.backup.%s", dbPath, time.Now().Format("20060102_150405")) + + // Read original database + data, err := os.ReadFile(dbPath) + if err != nil { + return fmt.Errorf("failed to read database for backup: %v", err) + } + + // Write backup + if err := os.WriteFile(backupFile, data, 0600); err != nil { + return fmt.Errorf("failed to create database backup: %v", err) + } + + fmt.Printf("Created database backup at: %s\n", backupFile) + } + + // Add new columns for 2FA - one at a time for SQLite compatibility + if err := tx.Exec(`ALTER TABLE users ADD COLUMN two_factor_secret VARCHAR(32)`).Error; err != nil { + return err + } + if err := tx.Exec(`ALTER TABLE users ADD COLUMN two_factor_enabled BOOLEAN DEFAULT FALSE`).Error; err != nil { + return err + } + if err := tx.Exec(`ALTER TABLE users ADD COLUMN backup_codes TEXT`).Error; err != nil { + return err + } + return nil + }, + Rollback: func(tx *gorm.DB) error { + // Remove 2FA columns - one at a time for SQLite compatibility + if err := tx.Exec(`ALTER TABLE users DROP COLUMN two_factor_secret`).Error; err != nil { + return err + } + if err := tx.Exec(`ALTER TABLE users DROP COLUMN two_factor_enabled`).Error; err != nil { + return err + } + if err := tx.Exec(`ALTER TABLE users DROP COLUMN backup_codes`).Error; err != nil { + return err + } + return nil + }, + } +} diff --git a/internal/db/migrations/migrations.go b/internal/db/migrations/migrations.go index 7b642ef..3a231aa 100644 --- a/internal/db/migrations/migrations.go +++ b/internal/db/migrations/migrations.go @@ -10,6 +10,7 @@ func InitMigrations(db *gorm.DB) *gormigrate.Gormigrate { migrations := []*gormigrate.Migration{ InitialSchema(), UpdateGDriveType(), + Add2FA(), } return gormigrate.New(db, gormigrate.DefaultOptions, migrations) diff --git a/internal/web/handlers/auth_handlers.go b/internal/web/handlers/auth_handlers.go index 4544703..20fda6e 100644 --- a/internal/web/handlers/auth_handlers.go +++ b/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") } diff --git a/internal/web/handlers/routes.go b/internal/web/handlers/routes.go index e57ca11..145e231 100644 --- a/internal/web/handlers/routes.go +++ b/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) diff --git a/internal/web/handlers/two_factor_handlers.go b/internal/web/handlers/two_factor_handlers.go new file mode 100644 index 0000000..bb67069 --- /dev/null +++ b/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(``)) + 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(``)) +}