Merge pull request #105 from PUFA-Computing/feat/2fa

feat: 2 Factor Authentication

Author: irfansaf

api/routes.go

@@ -85,7 +85,10 @@ func SetupRoutes() *gin.Engine {
 		userRoutes.PUT("/edit", userHandlers.EditUser)
 		userRoutes.DELETE("/delete", userHandlers.DeleteUser)
 		userRoutes.POST("/upload-profile-picture", userHandlers.UploadProfilePicture)
-		userRoutes.PUT("/:userID/update-user", userHandlers.AdminUpdateRoleAndStudentIDVerified)
+		userRoutes.PUT("/update-user", userHandlers.AdminUpdateRoleAndStudentIDVerified)
+		userRoutes.POST("/2fa/enable", userHandlers.EnableTwoFA)
+		userRoutes.POST("/2fa/verify", userHandlers.VerifyTwoFA)
+		userRoutes.POST("/2fa/disable", userHandlers.DisableTwoFA)
 
 		// ListEventsRegisteredByUser
 		userRoutes.GET("/registered-events", eventHandlers.ListEventsRegisteredByUser)

go.mod

@@ -38,6 +38,7 @@ require (
 	github.com/aws/aws-sdk-go-v2/service/sts v1.26.5 // indirect
 	github.com/aws/smithy-go v1.19.0 // indirect
 	github.com/benbjohnson/clock v1.3.0 // indirect
+	github.com/boombuler/barcode v1.0.1-0.20190219062509-6c824513bacc // indirect
 	github.com/bytedance/sonic v1.10.1 // indirect
 	github.com/cespare/xxhash/v2 v2.2.0 // indirect
 	github.com/chenzhuoyu/base64x v0.0.0-20230717121745-296ad89f973d // indirect
@@ -69,6 +70,7 @@ require (
 	github.com/pelletier/go-toml/v2 v2.1.0 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
+	github.com/pquerna/otp v1.4.0 // indirect
 	github.com/rogpeppe/go-internal v1.11.0 // indirect
 	github.com/stretchr/objx v0.5.0 // indirect
 	github.com/twitchyliquid64/golang-asm v0.15.1 // indirect

go.sum

@@ -44,6 +44,8 @@ github.com/aws/smithy-go v1.19.0 h1:KWFKQV80DpP3vJrrA9sVAHQ5gc2z8i4EzrLhLlWXcBM=
 github.com/aws/smithy-go v1.19.0/go.mod h1:NukqUGpCZIILqqiV0NIjeFh24kd/FAa4beRb6nbIUPE=
 github.com/benbjohnson/clock v1.3.0 h1:ip6w0uFQkncKQ979AypyG0ER7mqUSBdKLOgAle/AT8A=
 github.com/benbjohnson/clock v1.3.0/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZxNJlLklBHA=
+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/bsm/ginkgo/v2 v2.12.0 h1:Ny8MWAHyOepLGlLKYmXG4IEkioBysk6GpaRTLC8zwWs=
 github.com/bsm/ginkgo/v2 v2.12.0/go.mod h1:SwYbGRRDovPVboqFv0tPTcG1sN61LM1Z4ARdbAV9g4c=
 github.com/bsm/gomega v1.27.10 h1:yeMWxP2pV2fG3FgAODIY8EiRE3dy0aeFYt4l7wh6yKA=
@@ -183,6 +185,8 @@ github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
 github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 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/redis/go-redis/v9 v9.3.0 h1:RiVDjmig62jIWp7Kk4XVLs0hzV6pI3PyTnnL0cnn0u0=
 github.com/redis/go-redis/v9 v9.3.0/go.mod h1:hdY0cQFCN4fnSYT6TkisLufl/4W5UIXyv0b/CLO2V2M=
 github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/fJaraNFVN+nFs=

internal/database/app/user.go

@@ -8,6 +8,8 @@ import (
 	"errors"
 	"github.com/google/uuid"
 	"log"
+	"strconv"
+	"strings"
 )
 
 func GetUserByUsernameOrEmail(username string) (*models.User, error) {
@@ -83,7 +85,7 @@ func GetUserByEmail(email string) (*models.User, error) {
 	err := database.DB.QueryRow(context.Background(), query, email).Scan(
 		&userID, &user.Username, &user.Password, &user.FirstName, &user.MiddleName, &user.LastName, &user.Email,
 		&user.StudentID, &user.Major, &user.ProfilePicture, &user.DateOfBirth, &user.RoleID, &user.CreatedAt,
-		&user.UpdatedAt, &user.Year, &user.InstitutionName, &user.Gender)
+		&user.UpdatedAt, &user.Year, &user.InstitutionName, &user.Gender, &user.TwoFAEnabled, &user.TwoFAImage, &user.TwoFASecret)
 	if err != nil {
 		if errors.Is(err, sql.ErrNoRows) {
 			return nil, nil // User not found, return nil
@@ -101,13 +103,16 @@ func GetUserByID(userID uuid.UUID) (*models.User, error) {
 	var user models.User
 
 	err := database.DB.QueryRow(context.Background(), `
-		SELECT id, username, password, first_name, middle_name, last_name, email, student_id, major, profile_picture, date_of_birth, role_id, created_at, updated_at, year, institution_name, gender
+		SELECT id, username, password, first_name, middle_name, last_name, email, student_id, major, profile_picture, date_of_birth, role_id, created_at, updated_at, year, institution_name, gender, twofa_enabled, twofa_image, twofa_secret
 		FROM users WHERE id = $1`, userID).Scan(
 		&user.ID, &user.Username, &user.Password, &user.FirstName, &user.MiddleName, &user.LastName, &user.Email,
 		&user.StudentID, &user.Major, &user.ProfilePicture, &user.DateOfBirth, &user.RoleID, &user.CreatedAt,
-		&user.UpdatedAt, &user.Year, &user.InstitutionName, &user.Gender)
+		&user.UpdatedAt, &user.Year, &user.InstitutionName, &user.Gender, &user.TwoFAEnabled, &user.TwoFAImage, &user.TwoFASecret)
 
 	if err != nil {
+		if errors.Is(err, sql.ErrNoRows) {
+			return nil, nil
+		}
 		return nil, err
 	}
 
@@ -153,13 +158,104 @@ func CheckStudentIDExists(studentID string) (bool, error) {
 }
 
 func UpdateUser(UserID uuid.UUID, updatedUser *models.User) error {
-	_, err := database.DB.Exec(context.Background(), `
-		UPDATE users SET username = $1, password = $2, first_name = $3, middle_name = $4, last_name = $5, email = $6,
-		student_id = $7, major = $8, year = $9, role_id = $10, updated_at = $11, institution_name= $12, gender = $13
-		WHERE id = $14`,
-		updatedUser.Username, updatedUser.Password, updatedUser.FirstName, updatedUser.MiddleName, updatedUser.LastName,
-		updatedUser.Email, updatedUser.StudentID, updatedUser.Major, updatedUser.Year, updatedUser.RoleID,
-		updatedUser.UpdatedAt, updatedUser.InstitutionName, updatedUser.Gender, UserID)
+	query := "UPDATE users SET "
+	args := []interface{}{}
+	argID := 1
+
+	if updatedUser.Username != "" {
+		query += "username = $" + strconv.Itoa(argID) + ", "
+		args = append(args, updatedUser.Username)
+		argID++
+	}
+	if updatedUser.Password != "" {
+		query += "password = $" + strconv.Itoa(argID) + ", "
+		args = append(args, updatedUser.Password)
+		argID++
+	}
+	if updatedUser.FirstName != "" {
+		query += "first_name = $" + strconv.Itoa(argID) + ", "
+		args = append(args, updatedUser.FirstName)
+		argID++
+	}
+	if updatedUser.MiddleName != nil && *updatedUser.MiddleName != "" {
+		query += "middle_name = $" + strconv.Itoa(argID) + ", "
+		args = append(args, updatedUser.MiddleName)
+		argID++
+	}
+	if updatedUser.LastName != "" {
+		query += "last_name = $" + strconv.Itoa(argID) + ", "
+		args = append(args, updatedUser.LastName)
+		argID++
+	}
+	if updatedUser.Email != "" {
+		query += "email = $" + strconv.Itoa(argID) + ", "
+		args = append(args, updatedUser.Email)
+		argID++
+	}
+	if updatedUser.StudentID != "" {
+		query += "student_id = $" + strconv.Itoa(argID) + ", "
+		args = append(args, updatedUser.StudentID)
+		argID++
+	}
+	if updatedUser.Major != "" {
+		query += "major = $" + strconv.Itoa(argID) + ", "
+		args = append(args, updatedUser.Major)
+		argID++
+	}
+	if updatedUser.Year != "" {
+		query += "year = $" + strconv.Itoa(argID) + ", "
+		args = append(args, updatedUser.Year)
+		argID++
+	}
+	if updatedUser.DateOfBirth != nil {
+		query += "date_of_birth = $" + strconv.Itoa(argID) + ", "
+		args = append(args, updatedUser.DateOfBirth)
+		argID++
+	}
+	if updatedUser.RoleID != 0 {
+		query += "role_id = $" + strconv.Itoa(argID) + ", "
+		args = append(args, updatedUser.RoleID)
+		argID++
+	}
+	if !updatedUser.UpdatedAt.IsZero() {
+		query += "updated_at = $" + strconv.Itoa(argID) + ", "
+		args = append(args, updatedUser.UpdatedAt)
+		argID++
+	}
+	if updatedUser.InstitutionName != nil && *updatedUser.InstitutionName != "" {
+		query += "institution_name = $" + strconv.Itoa(argID) + ", "
+		args = append(args, updatedUser.InstitutionName)
+		argID++
+	}
+	if updatedUser.Gender != "" {
+		query += "gender = $" + strconv.Itoa(argID) + ", "
+		args = append(args, updatedUser.Gender)
+		argID++
+	}
+	if updatedUser.TwoFAEnabled != false {
+		query += "twofa_enabled = $" + strconv.Itoa(argID) + ", "
+		args = append(args, updatedUser.TwoFAEnabled)
+		argID++
+	}
+	if updatedUser.TwoFAImage != nil && *updatedUser.TwoFAImage != "" {
+		query += "twofa_image = $" + strconv.Itoa(argID) + ", "
+		args = append(args, updatedUser.TwoFAImage)
+		argID++
+	}
+	if updatedUser.TwoFASecret != nil && *updatedUser.TwoFASecret != "" {
+		query += "twofa_secret = $" + strconv.Itoa(argID) + ", "
+		args = append(args, updatedUser.TwoFASecret)
+		argID++
+	}
+
+	// Remove the last comma and space
+	query = strings.TrimSuffix(query, ", ")
+
+	// Add the WHERE clause
+	query += " WHERE id = $" + strconv.Itoa(argID)
+	args = append(args, UserID)
+
+	_, err := database.DB.Exec(context.Background(), query, args...)
 	return err
 }
 
@@ -195,6 +291,7 @@ func ListUsers() ([]models.User, error) {
 			&user.RoleID, &user.CreatedAt, &user.UpdatedAt, &user.Year, &user.EmailVerified,
 			&user.EmailVerificationToken, &user.PasswordResetToken, &user.PasswordResetExpires,
 			&user.StudentIDVerified, &user.StudentIDVerification, &user.InstitutionName, &user.Gender,
+			&user.TwoFAEnabled, &user.TwoFAImage, &user.TwoFASecret,
 		)
 		if err != nil {
 			log.Println("Error scanning row:", err)

internal/handlers/user/user_handlers.go

@@ -355,3 +355,68 @@ func (h *Handlers) UploadProfilePicture(c *gin.Context) {
 		"data":    user.ProfilePicture,
 	})
 }
+
+func (h *Handlers) EnableTwoFA(c *gin.Context) {
+	userID, err := (&auth.Handlers{}).ExtractUserIDAndCheckPermission(c, "users:2fa")
+	if err != nil {
+		c.JSON(http.StatusUnauthorized, gin.H{"success": false, "message": []string{err.Error()}})
+		return
+	}
+
+	qr, setupKey, err := h.UserService.EnableTwoFA(userID)
+	if err != nil {
+		c.JSON(http.StatusInternalServerError, gin.H{"success": false, "message": []string{err.Error()}})
+		return
+	}
+
+	c.JSON(http.StatusOK, gin.H{
+		"success": true,
+		"message": "Two Factor Authentication Enabled Successfully",
+		"data": gin.H{
+			"qr_image":  qr,
+			"setup_key": setupKey,
+		},
+	})
+}
+
+func (h *Handlers) VerifyTwoFA(c *gin.Context) {
+	userID, err := (&auth.Handlers{}).ExtractUserIDAndCheckPermission(c, "users:2fa")
+	if err != nil {
+		c.JSON(http.StatusUnauthorized, gin.H{"success": false, "message": []string{err.Error()}})
+		return
+	}
+
+	var request struct {
+		Code string `json:"code"`
+	}
+	if err := c.BindJSON(&request); err != nil {
+		c.JSON(http.StatusBadRequest, gin.H{"success": false, "message": err.Error()})
+		return
+	}
+	valid, err := h.UserService.VerifyTwoFA(userID, request.Code)
+	if err != nil {
+		c.JSON(http.StatusInternalServerError, gin.H{"success": false, "message": err.Error()})
+		return
+	}
+	if !valid {
+		c.JSON(http.StatusUnauthorized, gin.H{"success": false, "message": "Invalid TOTP Code"})
+		return
+	}
+	c.JSON(http.StatusOK, gin.H{"success": true, "message": "TOTP Verified Successfully"})
+}
+
+func (h *Handlers) DisableTwoFA(c *gin.Context) {
+	userID, err := (&auth.Handlers{}).ExtractUserIDAndCheckPermission(c, "users:2fa")
+	if err != nil {
+		c.JSON(http.StatusUnauthorized, gin.H{"success": false, "message": []string{err.Error()}})
+		return
+	}
+
+	err = h.UserService.DisableTwoFA(userID)
+	if err != nil {
+		c.JSON(http.StatusInternalServerError, gin.H{"success": false, "message": err.Error()})
+		return
+	}
+
+	c.JSON(http.StatusOK, gin.H{"success": true, "message": "Two Factor Authentication Disabled Successfully"})
+}

internal/models/user.go

@@ -30,4 +30,7 @@ type User struct {
 	InstitutionName        *string    `json:"institution_name"`
 	Gender                 string     `json:"gender"`
 	AdditionalNotes        *string    `json:"additional_notes"`
+	TwoFAEnabled           bool       `json:"2fa_enabled"`
+	TwoFAImage             *string    `json:"2fa_image"`
+	TwoFASecret            *string    `json:"2fa_secret"`
 }

internal/services/user_services.go

@@ -3,9 +3,14 @@ package services
 import (
 	"Backend/internal/database/app"
 	"Backend/internal/models"
+	"Backend/pkg/utils"
 	"errors"
+	"fmt"
 	"github.com/google/uuid"
+	"github.com/pquerna/otp"
+	"github.com/pquerna/otp/totp"
 	"log"
+	"time"
 )
 
 type UserService struct{}
@@ -106,3 +111,91 @@ func (us *UserService) AdminUpdateRoleAndStudentIDVerified(userID uuid.UUID, rol
 func (us *UserService) UploadProfilePicture(userID uuid.UUID, profilePicture string) error {
 	return app.UploadProfilePicture(userID, profilePicture)
 }
+
+func (us *UserService) EnableTwoFA(userID uuid.UUID) (string, string, error) {
+	user, err := app.GetUserByID(userID)
+	if err != nil {
+		return "", "", err
+	}
+
+	log.Println("Generating TOTP key")
+
+	secret, err := utils.GenerateTOTPKey(user.Email)
+	if err != nil {
+		return "", "", err
+	}
+
+	log.Println("Generated TOTP secret:", secret.Secret())
+
+	log.Println("Generating QR code")
+
+	qr, err := utils.GenerateQRCodeBase64(secret)
+	if err != nil {
+		return "", "", err
+	}
+
+	log.Println("Updating user")
+
+	secretStr := secret.Secret()
+
+	user.TwoFASecret = &secretStr
+	user.TwoFAEnabled = true
+	user.TwoFAImage = &qr
+
+	log.Println("User updated with TOTP secret:", secretStr)
+
+	err = app.UpdateUser(userID, user)
+	if err != nil {
+		return "", "", err
+	}
+
+	return qr, secretStr, nil
+}
+
+func (us *UserService) VerifyTwoFA(userID uuid.UUID, code string) (bool, error) {
+	user, err := app.GetUserByID(userID)
+	if err != nil {
+		return false, err
+	}
+
+	if user.TwoFASecret == nil {
+		log.Println("No TOTP secret found for user")
+		return false, fmt.Errorf("no TOTP secret found for user")
+	}
+
+	log.Println("Verifying TOTP code with secret:", *user.TwoFASecret)
+	log.Println("TOTP code to verify:", code)
+
+	// Ensure the correct settings for TOTP validation
+	valid, err := totp.ValidateCustom(code, *user.TwoFASecret, time.Now(), totp.ValidateOpts{
+		Period:    30,
+		Skew:      1,
+		Digits:    otp.DigitsEight,
+		Algorithm: otp.AlgorithmSHA256,
+	})
+
+	if err != nil {
+		return false, err
+	}
+
+	log.Println("Is TOTP code valid?", valid)
+	return valid, nil
+}
+
+func (us *UserService) DisableTwoFA(userID uuid.UUID) error {
+	user, err := app.GetUserByID(userID)
+	if err != nil {
+		return err
+	}
+
+	user.TwoFASecret = nil
+	user.TwoFAEnabled = false
+	user.TwoFAImage = nil
+
+	err = app.UpdateUser(userID, user)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}