feat: allow changing to a shared drive

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
+}