From b46ec659fc84017412e01ca4cffdffb8c2e30205 Mon Sep 17 00:00:00 2001
From: Ryan Stafford <ryan@nanorex.net>
Date: Mon, 3 Jul 2023 21:53:18 -0400
Subject: [PATCH] totp support

---
 go.mod                 |  2 +-
 go.sum                 |  2 ++
 routes.go              | 52 ++++++++++++++++++++++++++----------------
 templates/login.html   | 10 +++++++-
 templates/sidebar.html |  4 ++--
 5 files changed, 46 insertions(+), 24 deletions(-)

diff --git a/go.mod b/go.mod
index 88f3104..1f3e062 100644
--- a/go.mod
+++ b/go.mod
@@ -10,7 +10,7 @@ require (
 	github.com/gorilla/sessions v1.2.1 // indirect
 	github.com/gorilla/websocket v1.4.2 // indirect
 	github.com/julienschmidt/httprouter v1.3.0 // indirect
-	github.com/rystaf/go-lemmy v0.0.0-20230623191350-f39e3c8bdcb5 // indirect
+	github.com/rystaf/go-lemmy v0.0.0-20230704005320-c4b010dd339b // indirect
 	github.com/yuin/goldmark v1.5.4 // indirect
 	go.elara.ws/go-lemmy v0.17.3 // indirect
 	golang.org/x/text v0.10.0 // indirect
diff --git a/go.sum b/go.sum
index 66acff2..fdda18e 100644
--- a/go.sum
+++ b/go.sum
@@ -29,6 +29,8 @@ github.com/rystaf/go-lemmy v0.0.0-20230623191111-7ff8c74b1935 h1:zmzUz6PGRB8yQTT
 github.com/rystaf/go-lemmy v0.0.0-20230623191111-7ff8c74b1935/go.mod h1:nRSkTD+ARAHXtqlSPdf5q3hjHLP1ALsS1m5D3o86o+4=
 github.com/rystaf/go-lemmy v0.0.0-20230623191350-f39e3c8bdcb5 h1:MoI87uid2KqpLdUMZGK2HBOuxJMnPOJaar/4Og2PshM=
 github.com/rystaf/go-lemmy v0.0.0-20230623191350-f39e3c8bdcb5/go.mod h1:nRSkTD+ARAHXtqlSPdf5q3hjHLP1ALsS1m5D3o86o+4=
+github.com/rystaf/go-lemmy v0.0.0-20230704005320-c4b010dd339b h1:6z+gOUUvKwKQfgqEbxXS229gjr5V3HYg9bYbL9VHFdQ=
+github.com/rystaf/go-lemmy v0.0.0-20230704005320-c4b010dd339b/go.mod h1:nRSkTD+ARAHXtqlSPdf5q3hjHLP1ALsS1m5D3o86o+4=
 github.com/yuin/goldmark v1.5.4 h1:2uY/xC0roWy8IBEGLgB1ywIoEJFGmRrX21YQcvGZzjU=
 github.com/yuin/goldmark v1.5.4/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY=
 go.elara.ws/go-lemmy v0.17.3 h1:644k23BS2xqKJHJ9cHd8eyt1INpb5myqsBQQL2chBiA=
diff --git a/routes.go b/routes.go
index 94253a8..dca64b9 100644
--- a/routes.go
+++ b/routes.go
@@ -510,6 +510,7 @@ func Settings(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
 func SignUpOrLogin(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
 	state, err := Initialize(ps.ByName("host"), r)
 	if err != nil {
+		fmt.Println(err)
 		Render(w, "index.html", state)
 		return
 	}
@@ -517,11 +518,19 @@ func SignUpOrLogin(w http.ResponseWriter, r *http.Request, ps httprouter.Params)
 	var username string
 	switch r.FormValue("submit") {
 	case "log in":
-		resp, err := state.Client.Login(context.Background(), types.Login{
+		login := types.Login{
 			UsernameOrEmail: r.FormValue("username"),
 			Password:        r.FormValue("password"),
-		})
+		}
+		if r.FormValue("totp") != "" {
+			login.Totp2faToken = types.NewOptional(r.FormValue("totp"))
+		}
+		resp, err := state.Client.Login(context.Background(), login)
 		if err != nil {
+			if strings.Contains(fmt.Sprintf("%v", err), "missing_totp_token") {
+				state.Op = "2fa"
+			}
+			fmt.Println(err)
 			state.Error = err
 			state.GetSite()
 			state.GetCaptcha()
@@ -577,13 +586,6 @@ func SignUpOrLogin(w http.ResponseWriter, r *http.Request, ps httprouter.Params)
 		}
 	}
 	if token != "" {
-		if err != nil {
-			state.Error = err
-			state.GetSite()
-			state.GetCaptcha()
-			Render(w, "login.html", state)
-			return
-		}
 		state.GetUser(username)
 		setCookie(w, state.Host, "jwt", token)
 		userid := strconv.Itoa(state.User.PersonView.Person.ID)
@@ -671,18 +673,28 @@ func UserOp(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
 		deleteCookie(w, state.Host, "jwt")
 		deleteCookie(w, state.Host, "user")
 	case "login":
-		resp, err := state.Client.Login(context.Background(), types.Login{
-			UsernameOrEmail: r.FormValue("user"),
-			Password:        r.FormValue("pass"),
-		})
-		if err != nil {
-			state.Status = http.StatusUnauthorized
+		login := types.Login{
+			UsernameOrEmail: r.FormValue("username"),
+			Password:        r.FormValue("password"),
 		}
-		if resp.JWT.IsValid() {
-			state.GetUser(r.FormValue("user"))
-			setCookie(w, state.Host, "jwt", resp.JWT.String())
-			userid := strconv.Itoa(state.User.PersonView.Person.ID)
-			setCookie(w, state.Host, "user", state.User.PersonView.Person.Name+":"+userid)
+		if r.FormValue("totp") != "" {
+			login.Totp2faToken = types.NewOptional(r.FormValue("totp"))
+		}
+		resp, err := state.Client.Login(context.Background(), login)
+		if err != nil {
+			if strings.Contains(fmt.Sprintf("%v", err), "missing_totp_token") {
+				state.Op = "2fa"
+				Render(w, "login.html", state)
+				return
+			}
+			state.Status = http.StatusUnauthorized
+		} else if resp.JWT.IsValid() {
+			state.GetUser(r.FormValue("username"))
+			if state.User != nil {
+				setCookie(w, state.Host, "jwt", resp.JWT.String())
+				userid := strconv.Itoa(state.User.PersonView.Person.ID)
+				setCookie(w, state.Host, "user", state.User.PersonView.Person.Name+":"+userid)
+			}
 		}
 	case "create_community":
 		state.GetSite()
diff --git a/templates/login.html b/templates/login.html
index 78cc14a..700d16e 100644
--- a/templates/login.html
+++ b/templates/login.html
@@ -35,6 +35,7 @@
       <div class="error">{{.Error}}</div>
 {{ end }}
     <div class="signup">
+{{ if ne .Op "2fa" }}
     <div>
       <h2>create a new account</h2>
       <form method="POST">
@@ -81,9 +82,10 @@
         <input type="submit" name="submit" value="sign up">
       </form>
     </div>
+{{ end }}
     <div>
       <h2>login</h2>
-      <form method="POST">
+      <form method="POST" action="/{{host .Host}}/login">
         <label>
           username
           <div><input required name="username" type="text"></div>
@@ -92,6 +94,12 @@
           password
           <div><input required name="password" type="password"></div>
         </label>
+        {{ if  eq .Op "2fa" }}
+        <label>
+          2fa code
+          <div><input required name="totp" type="text"></div>
+        </label>
+        {{ end }}
         <input type="submit" name="submit" value="log in">
       </form>
     </div>
diff --git a/templates/sidebar.html b/templates/sidebar.html
index ff2050a..8137b77 100644
--- a/templates/sidebar.html
+++ b/templates/sidebar.html
@@ -32,8 +32,8 @@
 
 {{ if not .Session -}}
       <form class="login" method="post">
-        <input name="user" type="text" placeholder="username" maxlength="20">
-        <input name="pass" type="password" placeholder="password">
+        <input name="username" type="text" placeholder="username" maxlength="20">
+        <input name="password" type="password" placeholder="password">
         <div>
           <input type="submit" name="op" value="login">
         </div>