From b675a36304fc55a5f0c69c4436ec8be0c258d02b Mon Sep 17 00:00:00 2001
From: Chun-Hung Tseng <henrybear327@gmail.com>
Date: Sat, 1 Jul 2023 10:08:07 +0200
Subject: [PATCH] Add a new configuration option: UserAgent

---
 common/config.go         | 5 +++++
 common/proton_manager.go | 3 ++-
 common/user.go           | 3 ++-
 3 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/common/config.go b/common/config.go
index 07e6b3e..142e1b7 100644
--- a/common/config.go
+++ b/common/config.go
@@ -5,6 +5,7 @@ import "os"
 type Config struct {
 	/* Constants */
 	AppVersion string
+	UserAgent  string
 
 	/* Login */
 	FirstLoginCredential *FirstLoginCredentialData
@@ -37,6 +38,7 @@ type ReusableCredentialData struct {
 func NewConfigWithDefaultValues() *Config {
 	return &Config{
 		AppVersion: "",
+		UserAgent:  "",
 
 		FirstLoginCredential: &FirstLoginCredentialData{
 			Username: "",
@@ -62,6 +64,8 @@ func NewConfigWithDefaultValues() *Config {
 
 func NewConfigForIntegrationTests() *Config {
 	appVersion := os.Getenv("PROTON_API_BRIDGE_APP_VERSION")
+	userAgent := os.Getenv("PROTON_API_BRIDGE_USER_AGENT")
+
 	username := os.Getenv("PROTON_API_BRIDGE_TEST_USERNAME")
 	password := os.Getenv("PROTON_API_BRIDGE_TEST_PASSWORD")
 	twoFA := os.Getenv("PROTON_API_BRIDGE_TEST_TWOFA")
@@ -79,6 +83,7 @@ func NewConfigForIntegrationTests() *Config {
 
 	return &Config{
 		AppVersion: appVersion,
+		UserAgent:  userAgent,
 
 		FirstLoginCredential: &FirstLoginCredentialData{
 			Username: username,
diff --git a/common/proton_manager.go b/common/proton_manager.go
index 400b04a..bf9d70b 100644
--- a/common/proton_manager.go
+++ b/common/proton_manager.go
@@ -4,10 +4,11 @@ import (
 	"github.com/henrybear327/go-proton-api"
 )
 
-func getProtonManager(appVersion string) *proton.Manager {
+func getProtonManager(appVersion string, userAgent string) *proton.Manager {
 	/* Notes on API calls: if the app version is not specified, the api calls will be rejected. */
 	options := []proton.Option{
 		proton.WithAppVersion(appVersion),
+		proton.WithUserAgent(userAgent),
 	}
 	m := proton.New(options...)
 
diff --git a/common/user.go b/common/user.go
index 3273c1e..98b0785 100644
--- a/common/user.go
+++ b/common/user.go
@@ -48,7 +48,7 @@ func Login(ctx context.Context, config *Config) (*proton.Manager, *proton.Client
 	var addr []proton.Address
 
 	// get manager
-	m := getProtonManager(config.AppVersion)
+	m := getProtonManager(config.AppVersion, config.UserAgent)
 
 	if config.UseReusableLogin {
 		c = m.NewClient(config.ReusableCredential.UID, config.ReusableCredential.AccessToken, config.ReusableCredential.RefreshToken)
@@ -81,6 +81,7 @@ func Login(ctx context.Context, config *Config) (*proton.Manager, *proton.Client
 		if err != nil {
 			return nil, nil, nil, nil, nil, err
 		}
+		// log.Printf("Available scopes %#v", auth.Scope)
 
 		if auth.TwoFA.Enabled&proton.HasTOTP != 0 {
 			if config.FirstLoginCredential.TwoFA != "" {