Proton-API-Bridge/common/user.go

package common

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"log"
	"os"

	"github.com/ProtonMail/gopenpgp/v2/crypto"
	"github.com/ProtonMail/go-proton-api"
)

type ProtonDriveCredential struct {
	UID           string
	AccessToken   string
	RefreshToken  string
	SaltedKeyPass string
}

func cacheCredentialToFile(config *Config) error {
	if config.CredentialCacheFile != "" {
		str, err := json.Marshal(config.ReusableCredential)
		if err != nil {
			return err
		}

		file, err := os.Create(config.CredentialCacheFile)
		if err != nil {
			return err
		}
		defer file.Close()
		_, err = file.WriteString(string(str))
		if err != nil {
			return err
		}
	}

	return nil
}

/*
Log in methods
- username and password to log in
- UID and refresh token

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, authHandler proton.AuthHandler, deAuthHandler proton.Handler) (*proton.Manager, *proton.Client, *ProtonDriveCredential, *crypto.KeyRing, map[string]*crypto.KeyRing, map[string]proton.Address, error) {
	var c *proton.Client
	var auth proton.Auth
	var userKR *crypto.KeyRing
	var addrKRs map[string]*crypto.KeyRing
	var addrs map[string]proton.Address

	// get manager
	m := getProtonManager(config.AppVersion, config.UserAgent)

	if config.UseReusableLogin {
		c = m.NewClient(config.ReusableCredential.UID, config.ReusableCredential.AccessToken, config.ReusableCredential.RefreshToken)
		c.AddAuthHandler(authHandler)
		c.AddDeauthHandler(deAuthHandler)

		err := cacheCredentialToFile(config)
		if err != nil {
			return nil, nil, nil, nil, nil, nil, err
		}

		SaltedKeyPassByteArr, err := base64.StdEncoding.DecodeString(config.ReusableCredential.SaltedKeyPass)
		if err != nil {
			return nil, nil, nil, nil, nil, nil, err
		}
		userKR, addrKRs, addrs, _, err = getAccountKRs(ctx, c, nil, SaltedKeyPassByteArr)
		if err != nil {
			return nil, nil, nil, nil, nil, nil, err
		}

		return m, c, nil, userKR, addrKRs, addrs, nil
	} else {
		username := config.FirstLoginCredential.Username
		password := config.FirstLoginCredential.Password
		if username == "" || password == "" {
			return nil, nil, nil, nil, nil, nil, ErrUsernameAndPasswordRequired
		}

		// perform login
		var err error
		c, auth, err = m.NewClientWithLogin(ctx, username, []byte(password))
		if err != nil {
			return nil, nil, nil, nil, nil, nil, err
		}
		c.AddAuthHandler(authHandler)
		c.AddDeauthHandler(deAuthHandler)

		if auth.TwoFA.Enabled&proton.HasTOTP != 0 {
			if config.FirstLoginCredential.TwoFA != "" {
				err := c.Auth2FA(ctx, proton.Auth2FAReq{
					TwoFactorCode: config.FirstLoginCredential.TwoFA,
				})
				if err != nil {
					return nil, nil, nil, nil, nil, nil, err
				}
			} else {
				return nil, nil, nil, nil, nil, nil, Err2FACodeRequired
			}
		}

		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, addrs, saltedKeyPassByteArr, err = getAccountKRs(ctx, c, keyPass, nil)
		if err != nil {
			return nil, nil, nil, nil, nil, nil, err
		}

		saltedKeyPass := base64.StdEncoding.EncodeToString(saltedKeyPassByteArr)
		config.ReusableCredential.UID = auth.UID
		config.ReusableCredential.AccessToken = auth.AccessToken
		config.ReusableCredential.RefreshToken = auth.RefreshToken
		config.ReusableCredential.SaltedKeyPass = saltedKeyPass

		err = cacheCredentialToFile(config)
		if err != nil {
			return nil, nil, nil, nil, nil, nil, err
		}

		return m, c, &ProtonDriveCredential{
			UID:           auth.UID,
			AccessToken:   auth.AccessToken,
			RefreshToken:  auth.RefreshToken,
			SaltedKeyPass: saltedKeyPass,
		}, userKR, addrKRs, addrs, nil
	}
}

func Logout(ctx context.Context, config *Config, m *proton.Manager, c *proton.Client, userKR *crypto.KeyRing, addrKRs map[string]*crypto.KeyRing) error {
	defer m.Close()
	defer c.Close()

	if config.CredentialCacheFile == "" {
		log.Println("Logging out user")

		// log out
		err := c.AuthDelete(ctx)
		if err != nil {
			return err
		}

		// clear keyrings
		userKR.ClearPrivateParams()

		for i := range addrKRs {
			addrKRs[i].ClearPrivateParams()
		}
	}

	return nil
}