hailey.at
1package server
2
3import (
4 "bytes"
5 "encoding/base64"
6 "fmt"
7 "image/png"
8
9 "github.com/haileyok/cocoon/internal/helpers"
10 "github.com/haileyok/cocoon/models"
11 "github.com/labstack/echo/v4"
12 "github.com/pquerna/otp/totp"
13)
14
15func (s *Server) handleAccountTotpEnrollGet(e echo.Context) error {
16 urepo, sess, err := s.getSessionRepoOrErr(e)
17 if err != nil {
18 return e.Redirect(303, "/account/signin")
19 }
20
21 if urepo.TwoFactorType == models.TwoFactorTypeTotp {
22 sess.AddFlash("You have already enabled TOTP", "error")
23 sess.Save(e.Request(), e.Response())
24 return e.Redirect(303, "/account")
25 } else if urepo.TwoFactorType != models.TwoFactorTypeNone {
26 sess.AddFlash("You have already have another 2FA method enabled", "error")
27 sess.Save(e.Request(), e.Response())
28 return e.Redirect(303, "/account")
29 }
30
31 secret, err := totp.Generate(totp.GenerateOpts{
32 Issuer: s.config.Hostname,
33 AccountName: urepo.Repo.Did,
34 })
35 if err != nil {
36 s.logger.Error("error generating totp secret", "error", err)
37 return helpers.ServerError(e, nil)
38 }
39
40 sess.Values["totp-secret"] = secret.String()
41 if err := sess.Save(e.Request(), e.Response()); err != nil {
42 s.logger.Error("error saving session", "error", err)
43 return helpers.ServerError(e, nil)
44 }
45
46 var buf bytes.Buffer
47 img, err := secret.Image(200, 200)
48 if err != nil {
49 s.logger.Error("error generating image from secret", "error", err)
50 return helpers.ServerError(e, nil)
51 }
52 png.Encode(&buf, img)
53
54 b64img := fmt.Sprintf("data:image/png;base64,%s", base64.StdEncoding.EncodeToString(buf.Bytes()))
55
56 return e.Render(200, "totp_enroll.html", map[string]any{
57 "flashes": getFlashesFromSession(e, sess),
58 "Image": b64img,
59 })
60}
61
62type TotpEnrollRequest struct {
63 Code string `form:"code"`
64}
65
66func (s *Server) handleAccountTotpEnrollPost(e echo.Context) error {
67 urepo, sess, err := s.getSessionRepoOrErr(e)
68 if err != nil {
69 return e.Redirect(303, "/account/signin")
70 }
71
72 var req TotpEnrollRequest
73 if err := e.Bind(&req); err != nil {
74 s.logger.Error("error binding request for enroll totp", "error", err)
75 return helpers.ServerError(e, nil)
76 }
77
78 secret, ok := sess.Values["totp-secret"].(string)
79 if !ok {
80 return helpers.InputError(e, nil)
81 }
82
83 if !totp.Validate(req.Code, secret) {
84 sess.AddFlash("The provided code was not valid.", "error")
85 sess.Save(e.Request(), e.Response())
86 return e.Redirect(303, "/account/totp-enroll")
87 }
88
89 if err := s.db.Exec("UPDATE repos SET two_factor_type = ?, totp_secret = ? WHERE did = ?", nil, models.TwoFactorTypeTotp, secret, urepo.Repo.Did).Error; err != nil {
90 s.logger.Error("error updating database with totp token", "error", err)
91 return helpers.ServerError(e, nil)
92 }
93
94 sess.AddFlash("You have successfully enrolled in TOTP!", "success")
95 delete(sess.Values, "totp-secret")
96 sess.Save(e.Request(), e.Response())
97
98 return e.Redirect(303, "/account")
99}