Rework login process & implement MFA

controllers/mfa.go

@@ -2,6 +2,7 @@ package controllers
 
 import (
 	"bytes"
+	"encoding/base32"
 	"encoding/base64"
 	"fmt"
 	"image/png"
@@ -44,6 +45,15 @@ func TotpEnrollPage(c *fiber.Ctx) error {
 		AccountName: user.Email,
 	}
 
+	existingSecret := sess.Get("totp-enroll-secret")
+	if existingSecret != nil {
+		var b32NoPadding = base32.StdEncoding.WithPadding(base32.NoPadding)
+		options.Secret, err = b32NoPadding.DecodeString(existingSecret.(string))
+		if err != nil {
+			return err
+		}
+	}
+
 	key, err := totp.Generate(options)
 	if err != nil {
 		return err
@@ -65,6 +75,44 @@ func TotpEnrollPage(c *fiber.Ctx) error {
 		base64.StdEncoding.EncodeToString(buf.Bytes()),
 	)
 
+	var mfaError string
+	if c.Method() == "POST" {
+		otp := c.FormValue("otp")
+		if totp.Validate(otp, key.Secret()) {
+			err = user.TotpSercet.Scan(key.Secret())
+			if err != nil {
+				return err
+			}
+
+			result = db.Save(&user)
+			if result.Error != nil {
+				return err
+			}
+
+			sess.Set("totp-verified", "yes")
+			err = sess.Save()
+			if err != nil {
+				return err
+			}
+
+			redirectId := c.Query("redirect")
+			redirectUrl := "/"
+
+			if len(redirectId) > 0 {
+				redirectKey := fmt.Sprintf("redirect-%s", redirectId)
+				redirectVal := sess.Get(redirectKey)
+
+				if redirectVal != nil {
+					redirectUrl = redirectVal.(string)
+				}
+			}
+
+			return c.Redirect(redirectUrl)
+		} else {
+			mfaError = "Code temporaire invalide"
+		}
+	}
+
 	sess.Set("totp-enroll-secret", key.Secret())
 	err = sess.Save()
 	if err != nil {
@@ -72,8 +120,75 @@ func TotpEnrollPage(c *fiber.Ctx) error {
 	}
 
 	return c.Render("totp_enroll", fiber.Map{
-		"PageTitle": "Enregistrement multifacteur",
+		"PageTitle": "Enregistrement multifacteur (TOTP)",
 		"QrCode":    imgBase64,
 		"Secret":    key.Secret(),
+		"MfaError":  mfaError,
+	})
+}
+
+func TotpVerifyPage(c *fiber.Ctx) error {
+	db, err := helpers.GetDatabase()
+	if err != nil {
+		return err
+	}
+
+	sess, err := helpers.GetSessionStore(c)
+	if err != nil {
+		return err
+	}
+
+	totpVerified := sess.Get("totp-verified")
+	if totpVerified != nil {
+		return fiber.NewError(fiber.StatusForbidden, "Forbidden")
+	}
+
+	userid, err := helpers.GetSessionUserId(c)
+	if err != nil {
+		return err
+	}
+
+	var user models.User
+	result := db.First(&user, "id = ?", userid)
+
+	if result.Error != nil {
+		return err
+	}
+
+	if !user.TotpSercet.Valid {
+		return fiber.NewError(fiber.StatusForbidden, "Forbidden")
+	}
+
+	var mfaError string
+	if c.Method() == "POST" {
+		otp := c.FormValue("otp")
+		if totp.Validate(otp, user.TotpSercet.String) {
+			sess.Set("totp-verified", "yes")
+			err = sess.Save()
+			if err != nil {
+				return err
+			}
+
+			redirectId := c.Query("redirect")
+			redirectUrl := "/"
+
+			if len(redirectId) > 0 {
+				redirectKey := fmt.Sprintf("redirect-%s", redirectId)
+				redirectVal := sess.Get(redirectKey)
+
+				if redirectVal != nil {
+					redirectUrl = redirectVal.(string)
+				}
+			}
+
+			return c.Redirect(redirectUrl)
+		} else {
+			mfaError = "Code temporaire invalide"
+		}
+	}
+
+	return c.Render("totp_verify", fiber.Map{
+		"PageTitle": "Vérification multifacteur (TOTP)",
+		"MfaError":  mfaError,
 	})
 }