From d5bf0418cad1fbd0f27e8ff3d57cfd3d77b291f6 Mon Sep 17 00:00:00 2001
From: Chun-Hung Tseng <henrybear327@gmail.com>
Date: Wed, 6 Sep 2023 02:54:59 +0200
Subject: [PATCH] Add two-password mode

---
 common/config.go | 21 ++++++++++++---------
 common/error.go  |  1 +
 common/user.go   | 13 ++++++++++++-
 3 files changed, 25 insertions(+), 10 deletions(-)

diff --git a/common/config.go b/common/config.go
index a92bde6..46712fb 100644
--- a/common/config.go
+++ b/common/config.go
@@ -30,9 +30,10 @@ type Config struct {
 }
 
 type FirstLoginCredentialData struct {
-	Username string
-	Password string
-	TwoFA    string
+	Username        string
+	Password        string
+	MailboxPassword string
+	TwoFA           string
 }
 
 type ReusableCredentialData struct {
@@ -50,9 +51,10 @@ func NewConfigWithDefaultValues() *Config {
 		UserAgent:  "",
 
 		FirstLoginCredential: &FirstLoginCredentialData{
-			Username: "",
-			Password: "",
-			TwoFA:    "",
+			Username:        "",
+			Password:        "",
+			MailboxPassword: "",
+			TwoFA:           "",
 		},
 		ReusableCredential: &ReusableCredentialData{
 			UID:           "",
@@ -100,9 +102,10 @@ func NewConfigForIntegrationTests() *Config {
 		UserAgent:  userAgent,
 
 		FirstLoginCredential: &FirstLoginCredentialData{
-			Username: username,
-			Password: password,
-			TwoFA:    twoFA,
+			Username:        username,
+			Password:        password,
+			MailboxPassword: "",
+			TwoFA:           twoFA,
 		},
 		ReusableCredential: &ReusableCredentialData{
 			UID:           uid,
diff --git a/common/error.go b/common/error.go
index 6c3c3a3..dad2974 100644
--- a/common/error.go
+++ b/common/error.go
@@ -8,4 +8,5 @@ var (
 
 	ErrUsernameAndPasswordRequired = errors.New("username and password are required")
 	Err2FACodeRequired             = errors.New("this account requires a 2FA code")
+	ErrMailboxPasswordRequired     = errors.New("this account requires a mailbox password")
 )
diff --git a/common/user.go b/common/user.go
index 5e8496d..39c74da 100644
--- a/common/user.go
+++ b/common/user.go
@@ -106,9 +106,20 @@ func Login(ctx context.Context, config *Config, authHandler proton.AuthHandler,
 			}
 		}
 
+		var keyPass []byte
+		if auth.PasswordMode == proton.TwoPasswordMode {
+			if config.FirstLoginCredential.MailboxPassword != "" {
+				keyPass = []byte(config.FirstLoginCredential.MailboxPassword)
+			} else {
+				return nil, nil, nil, nil, nil, nil, ErrMailboxPasswordRequired
+			}
+		} else {
+			keyPass = []byte(config.FirstLoginCredential.Password)
+		}
+
 		// decrypt keyring
 		var saltedKeyPassByteArr []byte
-		userKR, addrKRs, addr, saltedKeyPassByteArr, err = getAccountKRs(ctx, c, []byte(password), nil)
+		userKR, addrKRs, addr, saltedKeyPassByteArr, err = getAccountKRs(ctx, c, keyPass, nil)
 		if err != nil {
 			return nil, nil, nil, nil, nil, nil, err
 		}