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
44 return helpers.ServerError(e, nil)
45 }
46
47 var buf bytes.Buffer
48 img, err := secret.Image(200, 200)
49 if err != nil {
50 s.logger.Error("error generating image from secret", "error", err)
51 return helpers.ServerError(e, nil)
52 }
53 png.Encode(&buf, img)
54
55 b64img := fmt.Sprintf("data:image/png;base64,%s", base64.StdEncoding.EncodeToString(buf.Bytes()))
56
57 return e.Render(200, "totp_enroll.html", map[string]any{
58 "flashes": getFlashesFromSession(e, sess),
59 "Image": b64img,
60 })
61}
62
63type TotpEnrollRequest struct {
64 Code string `form:"code"`
65}
66
67func (s *Server) handleAccountTotpEnrollPost(e echo.Context) error {
68 urepo, sess, err := s.getSessionRepoOrErr(e)
69 if err != nil {
70 return e.Redirect(303, "/account/signin")
71 }
72
73 var req TotpEnrollRequest
74 if err := e.Bind(&req); err != nil {
75 s.logger.Error("error binding request for enroll totp", "error", err)
76 return helpers.ServerError(e, nil)
77 }
78
79 secret, ok := sess.Values["totp-secret"].(string)
80 if !ok {
81 return helpers.InputError(e, nil)
82 }
83
84 if !totp.Validate(req.Code, secret) {
85 sess.AddFlash("The provided code was not valid.", "error")
86 sess.Save(e.Request(), e.Response())
87 return e.Redirect(303, "/account/totp-enroll")
88 }
89
90 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 {
91 s.logger.Error("error updating database with totp token", "error", err)
92 return helpers.ServerError(e, nil)
93 }
94
95 sess.AddFlash("You have successfully enrolled in TOTP!", "success")
96 delete(sess.Values, "totp-secret")
97 sess.Save(e.Request(), e.Response())
98
99 return e.Redirect(303, "/account")
100}