From 01d4a4ab086e140ecc3916cefd09e2009cffa741 Mon Sep 17 00:00:00 2001 From: Chun-Hung Tseng Date: Wed, 12 Jul 2023 21:46:29 +0200 Subject: [PATCH] Add authHandler and deAuthHandler to the ProtonDrive init function --- common/user.go | 9 +++++---- drive.go | 4 ++-- drive_test_helper.go | 2 +- 3 files changed, 8 insertions(+), 7 deletions(-) diff --git a/common/user.go b/common/user.go index 29fcf35..5e8496d 100644 --- a/common/user.go +++ b/common/user.go @@ -47,7 +47,7 @@ Log in methods Keyring decryption The password will be salted, and then used to decrypt the keyring. The salted password needs to be and can be cached, so the keyring can be re-decrypted when needed */ -func Login(ctx context.Context, config *Config) (*proton.Manager, *proton.Client, *ProtonDriveCredential, *crypto.KeyRing, map[string]*crypto.KeyRing, []proton.Address, error) { +func Login(ctx context.Context, config *Config, authHandler proton.AuthHandler, deAuthHandler proton.Handler) (*proton.Manager, *proton.Client, *ProtonDriveCredential, *crypto.KeyRing, map[string]*crypto.KeyRing, []proton.Address, error) { var c *proton.Client var auth proton.Auth var userKR *crypto.KeyRing @@ -59,8 +59,8 @@ func Login(ctx context.Context, config *Config) (*proton.Manager, *proton.Client if config.UseReusableLogin { c = m.NewClient(config.ReusableCredential.UID, config.ReusableCredential.AccessToken, config.ReusableCredential.RefreshToken) - - // TODO: register auth/deauth handler + c.AddAuthHandler(authHandler) + c.AddDeauthHandler(deAuthHandler) err := cacheCredentialToFile(config) if err != nil { @@ -90,7 +90,8 @@ func Login(ctx context.Context, config *Config) (*proton.Manager, *proton.Client if err != nil { return nil, nil, nil, nil, nil, nil, err } - // log.Printf("Available scopes %#v", auth.Scope) + c.AddAuthHandler(authHandler) + c.AddDeauthHandler(deAuthHandler) if auth.TwoFA.Enabled&proton.HasTOTP != 0 { if config.FirstLoginCredential.TwoFA != "" { diff --git a/drive.go b/drive.go index 317c697..52999b9 100644 --- a/drive.go +++ b/drive.go @@ -31,9 +31,9 @@ func NewDefaultConfig() *common.Config { return common.NewConfigWithDefaultValues() } -func NewProtonDrive(ctx context.Context, config *common.Config) (*ProtonDrive, *common.ProtonDriveCredential, error) { +func NewProtonDrive(ctx context.Context, config *common.Config, authHandler proton.AuthHandler, deAuthHandler proton.Handler) (*ProtonDrive, *common.ProtonDriveCredential, error) { /* Log in and logout */ - m, c, credentials, userKR, addrKRs, addrData, err := common.Login(ctx, config) + m, c, credentials, userKR, addrKRs, addrData, err := common.Login(ctx, config, authHandler, deAuthHandler) if err != nil { return nil, nil, err } diff --git a/drive_test_helper.go b/drive_test_helper.go index f0176ef..a5251c7 100644 --- a/drive_test_helper.go +++ b/drive_test_helper.go @@ -34,7 +34,7 @@ func setup(t *testing.T, replaceExistingDraft bool) (context.Context, context.Ca ctx, cancel := context.WithCancel(context.Background()) - protonDrive, auth, err := NewProtonDrive(ctx, config) + protonDrive, auth, err := NewProtonDrive(ctx, config, func(auth proton.Auth) {}, func() {}) if err != nil { t.Fatal(err) }