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 }