diff --git a/drive.go b/drive.go index f0765b5..7536ace 100644 --- a/drive.go +++ b/drive.go @@ -2,6 +2,7 @@ package proton_api_bridge import ( "context" + "fmt" "log" "github.com/henrybear327/Proton-API-Bridge/common" @@ -24,6 +25,7 @@ type ProtonDrive struct { m *proton.Manager userKR *crypto.KeyRing addrKRs map[string]*crypto.KeyRing + emailKRs map[string]*crypto.KeyRing addrData map[string]proton.Address signatureAddress string @@ -143,6 +145,7 @@ func NewProtonDrive(ctx context.Context, config *common.Config, authHandler prot m: m, userKR: userKR, addrKRs: addrKRs, + emailKRs: make(map[string]*crypto.KeyRing), addrData: addrData, signatureAddress: mainShare.Creator, @@ -194,6 +197,11 @@ func (protonDrive *ProtonDrive) getSignatureVerificationKeyring(emailAddresses [ return nil, err } } + if emailKR, ok := protonDrive.emailKRs[emailAddress]; ok { + if err := addKeysFromKR(ret, emailKR); err != nil { + return nil, err + } + } } if err := addKeysFromKR(ret, verificationAddrKRs...); err != nil { @@ -205,3 +213,23 @@ func (protonDrive *ProtonDrive) getSignatureVerificationKeyring(emailAddresses [ } return ret, nil } + +func (protonDrive *ProtonDrive) GetKeyRingForAddressID(addrID string) (*crypto.KeyRing, bool) { + kr, found := protonDrive.addrKRs[addrID] + return kr, found +} + +func (protonDrive *ProtonDrive) AddKeyRingForEmail(ctx context.Context, emailAddress string) (*crypto.KeyRing, error) { + pk, _, err := protonDrive.c.GetPublicKeys(ctx, emailAddress) + if err != nil { + return nil, fmt.Errorf("failed to get public keys for address %q: %w", emailAddress, err) + } + pkKR, err := pk.GetKeyRing() + if err != nil { + return nil, fmt.Errorf("failed to get keyring for address %q: %w", emailAddress, err) + } + + protonDrive.emailKRs[emailAddress] = pkKR + + return pkKR, nil +}