From 80e0fa6f3e0154338fb0af8a82ca32ae6281dd15 Mon Sep 17 00:00:00 2001
From: Simon Ser <contact@emersion.fr>
Date: Sun, 8 Dec 2019 12:14:41 +0100
Subject: [PATCH] Fix 2FA support

The 2FA information is now returned after the /auth request instead of
the /auth/info one.

Closes: https://github.com/emersion/hydroxide/issues/76
---
 auth/auth.go               |  4 ++--
 cmd/hydroxide/hydroxide.go | 27 +++++++++++++++++---------
 protonmail/auth.go         | 39 +++++++++++++++++++++++++++++---------
 protonmail/protonmail.go   |  2 ++
 4 files changed, 52 insertions(+), 20 deletions(-)

diff --git a/auth/auth.go b/auth/auth.go
index bf6c568..55d4b13 100644
--- a/auth/auth.go
+++ b/auth/auth.go
@@ -119,11 +119,11 @@ func authenticate(c *protonmail.Client, cachedAuth *CachedAuth, username string)
 			return nil, fmt.Errorf("cannot re-authenticate: failed to get auth info: %v", err)
 		}
 
-		if authInfo.TwoFactor == 1 {
+		if cachedAuth.TwoFactor.Enabled == 1 {
 			return nil, fmt.Errorf("cannot re-authenticate: two factor authentication enabled, please login manually")
 		}
 
-		auth, err = c.Auth(username, cachedAuth.LoginPassword, "", authInfo)
+		auth, err = c.Auth(username, cachedAuth.LoginPassword, authInfo)
 		if err != nil {
 			return nil, fmt.Errorf("cannot re-authenticate: %v", err)
 		}
diff --git a/cmd/hydroxide/hydroxide.go b/cmd/hydroxide/hydroxide.go
index dd97a91..45be460 100644
--- a/cmd/hydroxide/hydroxide.go
+++ b/cmd/hydroxide/hydroxide.go
@@ -153,18 +153,27 @@ func main() {
 				log.Fatal(err)
 			}
 
-			var twoFactorCode string
-			if authInfo.TwoFactor == 1 {
-				scanner := bufio.NewScanner(os.Stdin)
-				fmt.Printf("2FA code: ")
-				scanner.Scan()
-				twoFactorCode = scanner.Text()
-			}
-
-			a, err = c.Auth(username, loginPassword, twoFactorCode, authInfo)
+			a, err = c.Auth(username, loginPassword, authInfo)
 			if err != nil {
 				log.Fatal(err)
 			}
+
+			if a.TwoFactor.Enabled == 1 {
+				if a.TwoFactor.TOTP != 1 {
+					log.Fatal("Only TOTP is supported as a 2FA method")
+				}
+
+				scanner := bufio.NewScanner(os.Stdin)
+				fmt.Printf("2FA TOTP code: ")
+				scanner.Scan()
+				code := scanner.Text()
+
+				scope, err := c.AuthTOTP(code)
+				if err != nil {
+					log.Fatal(err)
+				}
+				a.Scope = scope
+			}
 		}
 
 		var mailboxPassword string
diff --git a/protonmail/auth.go b/protonmail/auth.go
index c241a39..b59e0c8 100644
--- a/protonmail/auth.go
+++ b/protonmail/auth.go
@@ -16,12 +16,6 @@ type authInfoReq struct {
 }
 
 type AuthInfo struct {
-	TwoFactor     int
-	TwoFactorInfo struct {
-		Enabled int
-		U2F     interface{} // TODO
-		TOTP    int
-	} `json:"2FA"`
 	version         int
 	modulus         string
 	serverEphemeral string
@@ -72,7 +66,6 @@ type authReq struct {
 	SRPSession      string
 	ClientEphemeral string
 	ClientProof     string
-	TwoFactorCode   string
 }
 
 type PasswordMode int
@@ -88,8 +81,14 @@ type Auth struct {
 	UID          string
 	AccessToken  string
 	RefreshToken string
+	UserID       string
 	EventID      string
 	PasswordMode PasswordMode
+	TwoFactor struct {
+		Enabled int
+		U2F interface{} // TODO
+		TOTP int
+	} `json:"2FA"`
 }
 
 type authResp struct {
@@ -106,7 +105,7 @@ func (resp *authResp) auth() *Auth {
 	return auth
 }
 
-func (c *Client) Auth(username, password, twoFactorCode string, info *AuthInfo) (*Auth, error) {
+func (c *Client) Auth(username, password string, info *AuthInfo) (*Auth, error) {
 	if info == nil {
 		var err error
 		if info, err = c.AuthInfo(username); err != nil {
@@ -124,7 +123,6 @@ func (c *Client) Auth(username, password, twoFactorCode string, info *AuthInfo)
 		SRPSession:      info.srpSession,
 		ClientEphemeral: base64.StdEncoding.EncodeToString(proofs.clientEphemeral),
 		ClientProof:     base64.StdEncoding.EncodeToString(proofs.clientProof),
-		TwoFactorCode:   twoFactorCode,
 	}
 
 	req, err := c.newJSONRequest(http.MethodPost, "/auth", reqData)
@@ -147,6 +145,29 @@ func (c *Client) Auth(username, password, twoFactorCode string, info *AuthInfo)
 	return auth, nil
 }
 
+func (c *Client) AuthTOTP(code string) (scope string, err error) {
+	reqData := struct {
+		TwoFactorCode string
+	}{
+		TwoFactorCode: code,
+	}
+
+	req, err := c.newJSONRequest(http.MethodPost, "/auth/2fa", reqData)
+	if err != nil {
+		return "", err
+	}
+
+	respData := struct {
+		resp
+		Scope string
+	}{}
+	if err := c.doJSON(req, &respData); err != nil {
+		return "", err
+	}
+
+	return respData.Scope, nil
+}
+
 type authRefreshReq struct {
 	UID          string `json:"Uid"`
 	RefreshToken string
diff --git a/protonmail/protonmail.go b/protonmail/protonmail.go
index 7db18a9..6f5cb85 100644
--- a/protonmail/protonmail.go
+++ b/protonmail/protonmail.go
@@ -157,6 +157,8 @@ func (c *Client) doJSON(req *http.Request, respData interface{}) error {
 		return err
 	}
 
+	//log.Printf("<< %v %v\n%#v", req.Method, req.URL.Path, respData)
+
 	if maybeError, ok := respData.(maybeError); ok {
 		if err := maybeError.Err(); err != nil {
 			log.Printf("request failed: %v %v: %v", req.Method, req.URL.String(), err)