go-urlsh/internal/web/setup-multifactor.go

package web

import (
	"bytes"
	"context"
	"encoding/base64"
	"encoding/json"
	"image/png"
	"log"
	"strconv"
	"time"

	"code.lila.network/adoralaura/go-urlsh/internal/constants"
	"code.lila.network/adoralaura/go-urlsh/internal/db"
	"code.lila.network/adoralaura/go-urlsh/internal/misc"
	"code.lila.network/adoralaura/go-urlsh/models"
	"github.com/gofiber/fiber/v2"
	"github.com/pquerna/otp/totp"
)

type mfaSetupResponse struct {
	Error          bool     `json:"error"`
	Message        string   `json:"message,omitempty"`
	RecoveryTokens []string `json:"recoverytokens,omitempty"`
}

func HandleAdminAccountMFASetupGet(c *fiber.Ctx) error {
	if !db.IsCookieValid(c.Cookies(constants.LoginCookieName, "")) {
		c.Location("/admin/")
		c.Status(fiber.StatusSeeOther)
		return nil
	}

	var mfaobject models.MFATemplateObject
	var mfaconfig models.MFAConfig
	mfaconfig.Active = false
	mfaconfig.ExpiresAt = time.Now().Add(15 * time.Minute)

	for i := 0; i < 6; i++ {
		mfaconfig.RecoveryCodes = append(mfaconfig.RecoveryCodes, misc.RandomString(8))
	}

	user, err := db.GetUserFromCookie(c.Cookies(constants.LoginCookieName))
	if err != nil {
		log.Println(err)
		fiber.NewError(fiber.StatusInternalServerError, "500 Internal Server Error")
	}
	mfaconfig.UserName = user.UserName

	key, err := totp.Generate(totp.GenerateOpts{
		Issuer:      "go-urlsh",
		AccountName: user.UserName,
	})
	if err != nil {
		log.Println(err)
		fiber.NewError(fiber.StatusInternalServerError, "500 Internal Server Error")
	}

	mfaconfig.TOTPSecret = key.Secret()
	mfaobject.Key = key.URL()

	var buf bytes.Buffer
	img, err := key.Image(200, 200)
	if err != nil {
		log.Println(err)
		fiber.NewError(fiber.StatusInternalServerError, "500 Internal Server Error")
	}

	png.Encode(&buf, img)
	base64img := base64.StdEncoding.EncodeToString(buf.Bytes())

	mfaobject.Image = base64img

	_, err = models.DB.NewInsert().Model(&mfaconfig).Exec(context.Background())
	if err != nil {
		log.Printf("[HandleAdminAccountMFASetupGet] Error inserting mfaconfig to DB: %q\n", err)
		fiber.NewError(fiber.StatusInternalServerError, "500 Internal Server Error")
	}

	misc.SetMFASetupCookie(c, strconv.Itoa(int(mfaconfig.ID)))
	return c.Render("setup-multifactor", mfaobject)
}

func HandleAdminAccountMFASetupPost(c *fiber.Ctx) error {
	if !db.IsCookieValid(c.Cookies(constants.LoginCookieName, "")) {
		c.Location("/admin/")
		c.Status(fiber.StatusSeeOther)
		return nil
	}

	var response mfaSetupResponse
	response.Error = true

	var token models.TokenRequest
	var config models.MFAConfig
	//var user models.User

	err := json.Unmarshal(c.Body(), &token)
	if err != nil {
		log.Println(err.Error())
		return fiber.NewError(fiber.StatusBadRequest, "400 Bad Request")
	}

	//if err := c.BodyParser(&token); err != nil {
	//	response.Message = "Token is invalid"
	//	c.Status(fiber.StatusBadRequest)
	//	c.JSON(response)
	//	return nil
	//}
	setupcookie := c.Cookies(constants.MFASetupCookieName, "")

	if setupcookie == "" {
		return fiber.NewError(fiber.StatusBadRequest, "400 Bad Request")
	}

	err = models.DB.NewSelect().Model(&config).Where("id = ?", setupcookie).Scan(context.Background())
	if err != nil {
		log.Printf("[HandleAdminAccountMFASetupGet] Error getting MFAConfig from DB: %q\n", err)
		return fiber.NewError(fiber.StatusInternalServerError, "500 Internal Server Error")
	}

	totpvalid := totp.Validate(token.Token, config.TOTPSecret)
	if totpvalid {
		response.Error = false
		response.Message = "Multifactor authentication was successfully set up!"
		response.RecoveryTokens = config.RecoveryCodes

		config.Active = true

		_, err := models.DB.NewUpdate().Model(&config).Column("active").WherePK().Exec(context.Background())
		if err != nil {
			log.Printf("[HandleAdminAccountMFASetupGet] Error getting MFAConfig from DB: %q\n", err)
			response.Message = "500: DB Error, see logs"
			c.Status(fiber.StatusInternalServerError)
			c.JSON(response)
			return nil
		}
		c.Status(fiber.StatusOK)
		c.JSON(response)
	} else {
		response.Error = true
		response.Message = "Token is invalid, please try again."
		c.Status(fiber.StatusBadRequest)
		c.JSON(response)
	}

	return nil
}