From 5ff13fac2ef6b8732af95917077cc4dd78dc5a47 Mon Sep 17 00:00:00 2001
From: rigormorrtiss <rigormorrtiss@gmail.com>
Date: Tue, 17 Aug 2021 19:03:55 +0400
Subject: [PATCH] feat: add support for MFA (2FA) login (#29)

* feat(ui): add ui.NewMfaLoginForm()

* feat: add support for MFA
---
 discordo.go | 53 ++++++++++++++++++++++++++++++++++++++---------------
 ui/forms.go | 25 ++++++++++++++++++++-----
 2 files changed, 58 insertions(+), 20 deletions(-)

diff --git a/discordo.go b/discordo.go
index e760504..dbbaf08 100644
--- a/discordo.go
+++ b/discordo.go
@@ -60,8 +60,7 @@ func main() {
 			SetRoot(mainFlex, true).
 			SetFocus(guildsTreeView)
 
-		discordSession = newSession("", "", t)
-		defer discordSession.Close()
+		discordSession = newSession(t)
 	} else {
 		loginForm = ui.NewLoginForm(onLoginFormLoginButtonSelected)
 		app.SetRoot(loginForm, true)
@@ -106,7 +105,7 @@ func onMessageInputFieldInputCapture(e *tcell.EventKey) *tcell.EventKey {
 	return e
 }
 
-func newSession(email string, password string, token string) (s *session.Session) {
+func newSession(token string) (s *session.Session) {
 	api.UserAgent = "" +
 		"Mozilla/5.0 (X11; Linux x86_64) " +
 		"AppleWebKit/537.36 (KHTML, like Gecko) " +
@@ -116,12 +115,7 @@ func newSession(email string, password string, token string) (s *session.Session
 	gateway.DefaultIdentity.Device = ""
 
 	var err error
-	if email != "" && password != "" {
-		s, err = session.Login(email, password, "")
-	} else if token != "" {
-		s, err = session.New(token)
-	}
-
+	s, err = session.New(token)
 	if err != nil {
 		panic(err)
 	}
@@ -260,12 +254,41 @@ func onLoginFormLoginButtonSelected() {
 		return
 	}
 
-	app.
-		SetRoot(mainFlex, true).
-		SetFocus(guildsTreeView)
+	// Make a scratch HTTP client without a token
+	client := api.NewClient("")
+	// Try to login without TOTP
+	l, err := client.Login(email, password)
+	if err != nil {
+		panic(err)
+	}
 
-	discordSession = newSession(email, password, "")
-	defer discordSession.Close()
+	if l.Token != "" && !l.MFA {
+		app.
+			SetRoot(mainFlex, true).
+			SetFocus(guildsTreeView)
 
-	go util.SetItem(kr, "token", discordSession.Token)
+		discordSession = newSession(l.Token)
+		go util.SetItem(kr, "token", l.Token)
+	} else if l.MFA {
+		loginForm = ui.NewMfaLoginForm(func() {
+			code := loginForm.GetFormItem(0).(*tview.InputField).GetText()
+			if code == "" {
+				return
+			}
+
+			l, err := client.TOTP(code, l.Ticket)
+			if err != nil {
+				panic(err)
+			}
+
+			app.
+				SetRoot(mainFlex, true).
+				SetFocus(guildsTreeView)
+
+			discordSession = newSession(l.Token)
+			go util.SetItem(kr, "token", l.Token)
+		})
+
+		app.SetRoot(loginForm, true)
+	}
 }
diff --git a/ui/forms.go b/ui/forms.go
index e0b6294..0692840 100644
--- a/ui/forms.go
+++ b/ui/forms.go
@@ -5,16 +5,31 @@ import (
 	"github.com/rivo/tview"
 )
 
-func NewLoginForm(onLoginFormLoginButtonSelected func()) (f *tview.Form) {
+func newBaseLoginForm() (f *tview.Form) {
 	f = tview.NewForm()
 	f.
-		AddInputField("Email", "", 0, nil, nil).
-		AddPasswordField("Password", "", 0, 0, nil).
-		AddButton("Login", onLoginFormLoginButtonSelected).
 		SetButtonsAlign(tview.AlignCenter).
 		SetButtonBackgroundColor(tcell.GetColor("#5865F2")).
 		SetBorder(true).
 		SetBorderPadding(0, 0, 1, 0)
 
-	return f
+	return
+}
+
+func NewLoginForm(onLoginFormLoginButtonSelected func()) (f *tview.Form) {
+	f = newBaseLoginForm()
+	f.
+		AddInputField("Email", "", 0, nil, nil).
+		AddPasswordField("Password", "", 0, 0, nil).
+		AddButton("Login", onLoginFormLoginButtonSelected)
+
+	return
+}
+
+func NewMfaLoginForm(onMfaLoginFormLoginButtonSelected func()) (f *tview.Form) {
+	f = newBaseLoginForm().
+		AddPasswordField("Code", "", 0, 0, nil).
+		AddButton("Login", onMfaLoginFormLoginButtonSelected)
+
+	return
 }