Merge pull request #40 from scelaris/pr2

protonmail: Decrypt private keys for all available addresses.

protonmail/auth.go

@@ -4,6 +4,7 @@ import (
 	"encoding/base64"
 	"errors"
 	"io/ioutil"
+	"log"
 	"net/http"
 	"strings"
 	"time"
@@ -193,6 +194,30 @@ func (c *Client) AuthRefresh(expiredAuth *Auth) (*Auth, error) {
 	return auth, nil
 }
 
+func unlockKey(e *openpgp.Entity, passphraseBytes []byte) error {
+	var privateKeys []*packet.PrivateKey
+
+	// e.PrivateKey is a signing key
+	if e.PrivateKey != nil {
+		privateKeys = append(privateKeys, e.PrivateKey)
+	}
+
+	// e.Subkeys are encryption keys
+	for _, subkey := range e.Subkeys {
+		if subkey.PrivateKey != nil {
+			privateKeys = append(privateKeys, subkey.PrivateKey)
+		}
+	}
+
+	for _, priv := range privateKeys {
+		if err := priv.Decrypt(passphraseBytes); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
 func (c *Client) Unlock(auth *Auth, passphrase string) (openpgp.EntityList, error) {
 	passphraseBytes := []byte(passphrase)
 	if auth.keySalt != "" {
@@ -218,26 +243,10 @@ func (c *Client) Unlock(auth *Auth, passphrase string) (openpgp.EntityList, erro
 	}
 
 	for _, e := range keyRing {
-		var privateKeys []*packet.PrivateKey
-
-		// e.PrivateKey is a signing key
-		if e.PrivateKey != nil {
-			privateKeys = append(privateKeys, e.PrivateKey)
-		}
-
-		// e.Subkeys are encryption keys
-		for _, subkey := range e.Subkeys {
-			if subkey.PrivateKey != nil {
-				privateKeys = append(privateKeys, subkey.PrivateKey)
-			}
-		}
-
-		for _, priv := range privateKeys {
-			if err := priv.Decrypt(passphraseBytes); err != nil {
+		if err := unlockKey(e, passphraseBytes); err != nil {
 			return nil, err
 		}
 	}
-	}
 
 	// Decrypt access token
 
@@ -260,6 +269,40 @@ func (c *Client) Unlock(auth *Auth, passphrase string) (openpgp.EntityList, erro
 	c.uid = auth.UID
 	c.accessToken = string(accessTokenBytes)
 	c.keyRing = keyRing
+
+	// Unlock additional private keys
+	addrs, err := c.ListAddresses()
+	if err != nil {
+		return nil, err
+	}
+
+	for _, addr := range addrs {
+		for _, key := range addr.Keys {
+			entity, err := key.Entity()
+			if err != nil {
+				return nil, err
+			}
+
+			found := false
+			for _, e := range keyRing {
+				if e.PrimaryKey.KeyIdString() == entity.PrimaryKey.KeyIdString() {
+					found = true
+					break
+				}
+			}
+			if found {
+				continue
+			}
+
+			if err := unlockKey(entity, passphraseBytes); err != nil {
+				log.Printf("failed to unlock key %v: %v", entity.PrimaryKey.KeyIdString(), err)
+				continue
+			}
+
+			keyRing = append(keyRing, entity)
+		}
+	}
+
 	return keyRing, nil
 }
 

protonmail/keys.go

@@ -9,13 +9,21 @@ import (
 	"golang.org/x/crypto/openpgp"
 )
 
+type PrivateKeyFlags int
+
+const (
+	PrivateKeyVerify  PrivateKeyFlags = 1
+	PrivateKeyEncrypt                 = 2
+)
+
 type PrivateKey struct {
 	ID          string
 	Version     int
-	PublicKey   string
+	Flags       PrivateKeyFlags
 	PrivateKey  string
 	Fingerprint string
 	Activation  interface{} // TODO
+	Primary     int
 }
 
 func (priv *PrivateKey) Entity() (*openpgp.Entity, error) {