[feat] Make EncryptionKey aware of the active address

6 files changed, 95 insertions, 78 deletions

diff --git a/src/leap/keymanager/__init__.py b/src/leap/keymanager/__init__.py
index 8a4efbe9..99ee1632 100644
--- a/src/leap/keymanager/__init__.py
+++ b/src/leap/keymanager/__init__.py
@@ -580,7 +580,7 @@ class KeyManager(object):
                 data, pubkey, passphrase, sign=signkey,
                 cipher_algo=cipher_algo)
             pubkey.encr_used = True
-            yield _keys.put_key(pubkey, address)
+            yield _keys.put_key(pubkey)
             defer.returnValue(encrypted)
 
         dpub = self.get_key(address, ktype, private=False,
@@ -637,7 +637,7 @@ class KeyManager(object):
                 signature = pubkey
                 if not pubkey.sign_used:
                     pubkey.sign_used = True
-                    yield _keys.put_key(pubkey, verify)
+                    yield _keys.put_key(pubkey)
                     defer.returnValue((decrypted, signature))
             else:
                 signature = InvalidSignature(
@@ -734,7 +734,7 @@ class KeyManager(object):
             if signed:
                 if not pubkey.sign_used:
                     pubkey.sign_used = True
-                    d = _keys.put_key(pubkey, address)
+                    d = _keys.put_key(pubkey)
                     d.addCallback(lambda _: pubkey)
                     return d
                 return pubkey
@@ -765,20 +765,16 @@ class KeyManager(object):
         _keys = self._wrapper_map[type(key)]
         return _keys.delete_key(key)
 
-    def put_key(self, key, address):
+    def put_key(self, key):
         """
         Put key bound to address in local storage.
 
         :param key: The key to be stored
         :type key: EncryptionKey
-        :param address: address for which this key will be active
-        :type address: str
 
         :return: A Deferred which fires when the key is in the storage, or
-                 which fails with KeyAddressMismatch if address doesn't match
-                 any uid on the key or fails with KeyNotValidUpdate if a key
-                 with the same uid exists and the new one is not a valid update
-                 for it.
+                 which fails with KeyNotValidUpdate if a key with the same
+                 uid exists and the new one is not a valid update for it.
         :rtype: Deferred
 
         :raise UnsupportedKeyTypeError: if invalid key type
@@ -787,11 +783,6 @@ class KeyManager(object):
         self._assert_supported_key_type(ktype)
         _keys = self._wrapper_map[ktype]
 
-        if address not in key.address:
-            return defer.fail(
-                KeyAddressMismatch("UID %s found, but expected %s"
-                                   % (str(key.address), address)))
-
         def old_key_not_found(failure):
             if failure.check(KeyNotFound):
                 return None
@@ -800,13 +791,13 @@ class KeyManager(object):
 
         def check_upgrade(old_key):
             if key.private or can_upgrade(key, old_key):
-                return _keys.put_key(key, address)
+                return _keys.put_key(key)
             else:
                 raise KeyNotValidUpgrade(
                     "Key %s can not be upgraded by new key %s"
                     % (old_key.fingerprint, key.fingerprint))
 
-        d = _keys.get_key(address, private=key.private)
+        d = _keys.get_key(key.address, private=key.private)
         d.addErrback(old_key_not_found)
         d.addCallback(check_upgrade)
         return d
@@ -838,11 +829,11 @@ class KeyManager(object):
         self._assert_supported_key_type(ktype)
         _keys = self._wrapper_map[ktype]
 
-        pubkey, privkey = _keys.parse_ascii_key(key)
+        pubkey, privkey = _keys.parse_ascii_key(key, address)
         pubkey.validation = validation
-        d = self.put_key(pubkey, address)
+        d = self.put_key(pubkey)
         if privkey is not None:
-            d.addCallback(lambda _: self.put_key(privkey, address))
+            d.addCallback(lambda _: self.put_key(privkey))
         return d
 
     @defer.inlineCallbacks
@@ -878,12 +869,12 @@ class KeyManager(object):
         ascii_content = yield self._get_with_combined_ca_bundle(uri)
 
         # XXX parse binary keys
-        pubkey, _ = _keys.parse_ascii_key(ascii_content)
+        pubkey, _ = _keys.parse_ascii_key(ascii_content, address)
         if pubkey is None:
             raise KeyNotFound(uri)
 
         pubkey.validation = validation
-        yield self.put_key(pubkey, address)
+        yield self.put_key(pubkey)
 
     def _assert_supported_key_type(self, ktype):
         """
diff --git a/src/leap/keymanager/keys.py b/src/leap/keymanager/keys.py
index 68e3fada..38d66b5f 100644
--- a/src/leap/keymanager/keys.py
+++ b/src/leap/keymanager/keys.py
@@ -46,6 +46,7 @@ logger = logging.getLogger(__name__)
 #
 
 KEY_VERSION_KEY = 'version'
+KEY_UIDS_KEY = 'uids'
 KEY_ADDRESS_KEY = 'address'
 KEY_TYPE_KEY = 'type'
 KEY_FINGERPRINT_KEY = 'fingerprint'
@@ -126,12 +127,14 @@ def build_key_from_dict(kClass, key, active=None):
     :return: An instance of the key.
     :rtype: C{kClass}
     """
+    address = None
     validation = ValidationLevels.Weak_Chain
     last_audited_at = None
     encr_used = False
     sign_used = False
 
     if active:
+        address = active[KEY_ADDRESS_KEY]
         try:
             validation = ValidationLevels.get(active[KEY_VALIDATION_KEY])
         except ValueError:
@@ -146,7 +149,8 @@ def build_key_from_dict(kClass, key, active=None):
     refreshed_at = _to_datetime(key[KEY_REFRESHED_AT_KEY])
 
     return kClass(
-        key[KEY_ADDRESS_KEY],
+        address=address,
+        uids=key[KEY_UIDS_KEY],
         fingerprint=key[KEY_FINGERPRINT_KEY],
         key_data=key[KEY_DATA_KEY],
         private=key[KEY_PRIVATE_KEY],
@@ -188,12 +192,15 @@ class EncryptionKey(object):
 
     __metaclass__ = ABCMeta
 
-    def __init__(self, address, fingerprint="",
+    def __init__(self, address=None, uids=[], fingerprint="",
                  key_data="", private=False, length=0, expiry_date=None,
                  validation=ValidationLevels.Weak_Chain, last_audited_at=None,
                  refreshed_at=None, encr_used=False, sign_used=False):
-        # TODO: it should know its own active address
         self.address = address
+        if not uids and address:
+            self.uids = [address]
+        else:
+            self.uids = uids
         self.fingerprint = fingerprint
         self.key_data = key_data
         self.private = private
@@ -217,7 +224,7 @@ class EncryptionKey(object):
         refreshed_at = _to_unix_time(self.refreshed_at)
 
         return json.dumps({
-            KEY_ADDRESS_KEY: self.address,
+            KEY_UIDS_KEY: self.uids,
             KEY_TYPE_KEY: self.__class__.__name__,
             KEY_FINGERPRINT_KEY: self.fingerprint,
             KEY_DATA_KEY: self.key_data,
@@ -229,7 +236,7 @@ class EncryptionKey(object):
             KEY_TAGS_KEY: [KEYMANAGER_KEY_TAG],
         })
 
-    def get_active_json(self, address):
+    def get_active_json(self):
         """
         Return a JSON string describing this key.
 
@@ -239,7 +246,7 @@ class EncryptionKey(object):
         last_audited_at = _to_unix_time(self.last_audited_at)
 
         return json.dumps({
-            KEY_ADDRESS_KEY: address,
+            KEY_ADDRESS_KEY: self.address,
             KEY_TYPE_KEY: self.__class__.__name__ + KEYMANAGER_ACTIVE_TYPE,
             KEY_FINGERPRINT_KEY: self.fingerprint,
             KEY_PRIVATE_KEY: self.private,
@@ -374,14 +381,12 @@ class EncryptionScheme(object):
         pass
 
     @abstractmethod
-    def put_key(self, key, address):
+    def put_key(self, key):
         """
         Put a key in local storage.
 
         :param key: The key to be stored.
         :type key: EncryptionKey
-        :param address: address for which this key will be active.
-        :type address: str
 
         :return: A Deferred which fires when the key is in the storage.
         :rtype: Deferred
@@ -496,7 +501,7 @@ class EncryptionScheme(object):
         :rtype: Deferred
         """
         def log_key_doc(doc):
-            logger.error("\t%s: %s" % (doc.content[KEY_ADDRESS_KEY],
+            logger.error("\t%s: %s" % (doc.content[KEY_UIDS_KEY],
                                        doc.content[KEY_FINGERPRINT_KEY]))
 
         def cmp_key(d1, d2):
diff --git a/src/leap/keymanager/migrator.py b/src/leap/keymanager/migrator.py
index 11cf2439..9e4ae77c 100644
--- a/src/leap/keymanager/migrator.py
+++ b/src/leap/keymanager/migrator.py
@@ -32,6 +32,8 @@ from leap.keymanager.keys import (
     KEYMANAGER_ACTIVE_TAG,
 
     KEYMANAGER_DOC_VERSION,
+    KEY_ADDRESS_KEY,
+    KEY_UIDS_KEY,
     KEY_VERSION_KEY,
     KEY_FINGERPRINT_KEY,
     KEY_VALIDATION_KEY,
@@ -164,6 +166,8 @@ class KeyDocumentsMigrator(object):
             return succeed(None)
 
         key.content[KEY_VERSION_KEY] = KEYMANAGER_DOC_VERSION
+        key.content[KEY_UIDS_KEY] = key.content[KEY_ADDRESS_KEY]
+        del key.content[KEY_ADDRESS_KEY]
         del key.content[KEY_ID_KEY]
         del key.content[KEY_VALIDATION_KEY]
         del key.content[KEY_LAST_AUDITED_AT_KEY]
diff --git a/src/leap/keymanager/openpgp.py b/src/leap/keymanager/openpgp.py
index 0f162969..0d5a8667 100644
--- a/src/leap/keymanager/openpgp.py
+++ b/src/leap/keymanager/openpgp.py
@@ -41,7 +41,7 @@ from leap.keymanager.keys import (
     build_key_from_dict,
     TYPE_FINGERPRINT_PRIVATE_INDEX,
     TYPE_ADDRESS_PRIVATE_INDEX,
-    KEY_ADDRESS_KEY,
+    KEY_UIDS_KEY,
     KEY_FINGERPRINT_KEY,
     KEYMANAGER_ACTIVE_TYPE,
 )
@@ -200,7 +200,7 @@ class OpenPGPKey(EncryptionKey):
     Base class for OpenPGP keys.
     """
 
-    def __init__(self, address, gpgbinary=None, **kwargs):
+    def __init__(self, address=None, gpgbinary=None, **kwargs):
         self._gpgbinary = gpgbinary
         super(OpenPGPKey, self).__init__(address, **kwargs)
 
@@ -215,7 +215,7 @@ class OpenPGPKey(EncryptionKey):
         with TempGPGWrapper(keys=[self], gpgbinary=self._gpgbinary) as gpg:
             res = gpg.list_sigs(self.fingerprint)
             for uid, sigs in res.sigs.iteritems():
-                if _parse_address(uid) in self.address:
+                if _parse_address(uid) in self.uids:
                     return sigs
 
         return []
@@ -335,8 +335,9 @@ class OpenPGPScheme(EncryptionScheme):
                     key = gpg.list_keys(secret=secret).pop()
                     openpgp_key = self._build_key_from_gpg(
                         key,
-                        gpg.export_keys(key['fingerprint'], secret=secret))
-                    d = self.put_key(openpgp_key, address)
+                        gpg.export_keys(key['fingerprint'], secret=secret),
+                        address)
+                    d = self.put_key(openpgp_key)
                     deferreds.append(d)
                 yield defer.gatherResults(deferreds)
 
@@ -368,10 +369,10 @@ class OpenPGPScheme(EncryptionScheme):
             if keydoc is None:
                 raise errors.KeyNotFound(address)
             leap_assert(
-                address in keydoc.content[KEY_ADDRESS_KEY],
+                address in keydoc.content[KEY_UIDS_KEY],
                 'Wrong address in key %s. Expected %s, found %s.'
                 % (keydoc.content[KEY_FINGERPRINT_KEY], address,
-                   keydoc.content[KEY_ADDRESS_KEY]))
+                   keydoc.content[KEY_UIDS_KEY]))
             key = build_key_from_dict(OpenPGPKey, keydoc.content,
                                       activedoc.content)
             key._gpgbinary = self._gpgbinary
@@ -381,13 +382,15 @@ class OpenPGPScheme(EncryptionScheme):
         d.addCallback(build_key)
         return d
 
-    def parse_ascii_key(self, key_data):
+    def parse_ascii_key(self, key_data, address=None):
         """
         Parses an ascii armored key (or key pair) data and returns
         the OpenPGPKey keys.
 
         :param key_data: the key data to be parsed.
         :type key_data: str or unicode
+        :param address: Active address for the key.
+        :type address: str
 
         :returns: the public key and private key (if applies) for that data.
         :rtype: (public, private) -> tuple(OpenPGPKey, OpenPGPKey)
@@ -408,12 +411,13 @@ class OpenPGPScheme(EncryptionScheme):
         openpgp_privkey = None
         if privkey:
             # build private key
-            openpgp_privkey = self._build_key_from_gpg(priv_info, privkey)
+            openpgp_privkey = self._build_key_from_gpg(priv_info, privkey,
+                                                       address)
             leap_check(pub_info['fingerprint'] == priv_info['fingerprint'],
                        'Fingerprints for public and private key differ.',
                        errors.KeyFingerprintMismatch)
         # build public key
-        openpgp_pubkey = self._build_key_from_gpg(pub_info, pubkey)
+        openpgp_pubkey = self._build_key_from_gpg(pub_info, pubkey, address)
 
         return (openpgp_pubkey, openpgp_privkey)
 
@@ -433,12 +437,13 @@ class OpenPGPScheme(EncryptionScheme):
 
         openpgp_privkey = None
         try:
-            openpgp_pubkey, openpgp_privkey = self.parse_ascii_key(key_data)
+            openpgp_pubkey, openpgp_privkey = self.parse_ascii_key(
+                key_data, address)
         except (errors.KeyAddressMismatch, errors.KeyFingerprintMismatch) as e:
             return defer.fail(e)
 
         def put_key(_, key):
-            return self.put_key(key, address)
+            return self.put_key(key)
 
         d = defer.succeed(None)
         if openpgp_pubkey is not None:
@@ -447,14 +452,12 @@ class OpenPGPScheme(EncryptionScheme):
             d.addCallback(put_key, openpgp_privkey)
         return d
 
-    def put_key(self, key, address):
+    def put_key(self, key):
         """
         Put C{key} in local storage.
 
         :param key: The key to be stored.
         :type key: OpenPGPKey
-        :param address: address for which this key will be active.
-        :type address: str
 
         :return: A Deferred which fires when the key is in the storage.
         :rtype: Deferred
@@ -471,31 +474,36 @@ class OpenPGPScheme(EncryptionScheme):
 
             key.merge(oldkey)
             keydoc.set_json(key.get_json())
-            deferred_key = self._soledad.put_doc(keydoc)
-
-            active_json = key.get_active_json(address)
-            if activedoc:
-                activedoc.set_json(active_json)
-                deferred_active = self._soledad.put_doc(activedoc)
-            else:
-                deferred_active = self._soledad.create_doc_from_json(
-                    active_json)
-
-            return defer.gatherResults([deferred_key, deferred_active])
+            d = self._soledad.put_doc(keydoc)
+            d.addCallback(put_active, activedoc)
+            return d
 
         def put_new_key(activedoc):
             deferreds = []
             if activedoc:
                 d = self._soledad.delete_doc(activedoc)
                 deferreds.append(d)
-            for json in [key.get_json(), key.get_active_json(address)]:
+            for json in [key.get_json(), key.get_active_json()]:
                 d = self._soledad.create_doc_from_json(json)
                 deferreds.append(d)
             return defer.gatherResults(deferreds)
 
-        dk = self._get_key_doc_from_fingerprint(key.fingerprint, key.private)
-        da = self._get_active_doc_from_address(address, key.private)
-        d = defer.gatherResults([dk, da])
+        def put_active(_, activedoc):
+            active_json = key.get_active_json()
+            if activedoc:
+                activedoc.set_json(active_json)
+                d = self._soledad.put_doc(activedoc)
+            else:
+                d = self._soledad.create_doc_from_json(active_json)
+            return d
+
+        def get_active_doc(keydoc):
+            d = self._get_active_doc_from_address(key.address, key.private)
+            d.addCallback(lambda activedoc: (keydoc, activedoc))
+            return d
+
+        d = self._get_key_doc_from_fingerprint(key.fingerprint, key.private)
+        d.addCallback(get_active_doc)
         d.addCallback(merge_and_put)
         return d
 
@@ -533,7 +541,7 @@ class OpenPGPScheme(EncryptionScheme):
         d.addCallback(get_key_from_active_doc)
         return d
 
-    def _build_key_from_gpg(self, key, key_data):
+    def _build_key_from_gpg(self, key, key_data, address=None):
         """
         Build an OpenPGPKey for C{address} based on C{key} from
         local gpg storage.
@@ -541,6 +549,8 @@ class OpenPGPScheme(EncryptionScheme):
         ASCII armored GPG key data has to be queried independently in this
         wrapper, so we receive it in C{key_data}.
 
+        :param address: Active address for the key.
+        :type address: str
         :param key: Key obtained from GPG storage.
         :type key: dict
         :param key_data: Key data obtained from GPG storage.
@@ -548,7 +558,7 @@ class OpenPGPScheme(EncryptionScheme):
         :return: An instance of the key.
         :rtype: OpenPGPKey
         """
-        return build_gpg_key(key, key_data, self._gpgbinary)
+        return build_gpg_key(key, key_data, address, self._gpgbinary)
 
     def delete_key(self, key):
         """
@@ -852,16 +862,20 @@ def process_ascii_key(key_data, gpgbinary, secret=False):
     return info, key
 
 
-def build_gpg_key(key_info, key_data, gpgbinary=None):
+def build_gpg_key(key_info, key_data, address=None, gpgbinary=None):
     expiry_date = None
     if key_info['expires']:
         expiry_date = datetime.fromtimestamp(int(key_info['expires']))
-    address = []
+    uids = []
     for uid in key_info['uids']:
-        address.append(_parse_address(uid))
+        uids.append(_parse_address(uid))
+    if address and address not in uids:
+        raise errors.KeyAddressMismatch("UIDs %s found, but expected %s"
+                                        % (str(uids), address))
 
     return OpenPGPKey(
-        address,
+        address=address,
+        uids=uids,
         gpgbinary=gpgbinary,
         fingerprint=key_info['fingerprint'],
         key_data=key_data,
diff --git a/src/leap/keymanager/tests/test_keymanager.py b/src/leap/keymanager/tests/test_keymanager.py
index 2fe9e4cd..6347d56e 100644
--- a/src/leap/keymanager/tests/test_keymanager.py
+++ b/src/leap/keymanager/tests/test_keymanager.py
@@ -76,7 +76,7 @@ class KeyManagerUtilTestCase(unittest.TestCase):
 
     def test_build_key_from_dict(self):
         kdict = {
-            'address': [ADDRESS],
+            'uids': [ADDRESS],
             'fingerprint': KEY_FINGERPRINT,
             'key_data': PUBLIC_KEY,
             'private': False,
@@ -94,7 +94,7 @@ class KeyManagerUtilTestCase(unittest.TestCase):
         }
         key = build_key_from_dict(OpenPGPKey, kdict, adict)
         self.assertEqual(
-            kdict['address'], key.address,
+            kdict['uids'], key.uids,
             'Wrong data in key.')
         self.assertEqual(
             kdict['fingerprint'], key.fingerprint,
@@ -118,6 +118,9 @@ class KeyManagerUtilTestCase(unittest.TestCase):
             datetime.fromtimestamp(kdict['refreshed_at']), key.refreshed_at,
             'Wrong data in key.')
         self.assertEqual(
+            adict['address'], key.address,
+            'Wrong data in key.')
+        self.assertEqual(
             ValidationLevels.get(adict['validation']), key.validation,
             'Wrong data in key.')
         self.assertEqual(
@@ -137,12 +140,12 @@ class KeyManagerKeyManagementTestCase(KeyManagerWithSoledadTestCase):
         # get public keys
         keys = yield km.get_all_keys(False)
         self.assertEqual(len(keys), 1, 'Wrong number of keys')
-        self.assertTrue(ADDRESS in keys[0].address)
+        self.assertTrue(ADDRESS in keys[0].uids)
         self.assertFalse(keys[0].private)
         # get private keys
         keys = yield km.get_all_keys(True)
         self.assertEqual(len(keys), 1, 'Wrong number of keys')
-        self.assertTrue(ADDRESS in keys[0].address)
+        self.assertTrue(ADDRESS in keys[0].uids)
         self.assertTrue(keys[0].private)
 
     @defer.inlineCallbacks
@@ -153,7 +156,7 @@ class KeyManagerKeyManagementTestCase(KeyManagerWithSoledadTestCase):
         key = yield km.get_key(ADDRESS, OpenPGPKey, private=False,
                                fetch_remote=False)
         self.assertTrue(key is not None)
-        self.assertTrue(ADDRESS in key.address)
+        self.assertTrue(ADDRESS in key.uids)
         self.assertEqual(
             key.fingerprint.lower(), KEY_FINGERPRINT.lower())
         self.assertFalse(key.private)
@@ -166,7 +169,7 @@ class KeyManagerKeyManagementTestCase(KeyManagerWithSoledadTestCase):
         key = yield km.get_key(ADDRESS, OpenPGPKey, private=True,
                                fetch_remote=False)
         self.assertTrue(key is not None)
-        self.assertTrue(ADDRESS in key.address)
+        self.assertTrue(ADDRESS in key.uids)
         self.assertEqual(
             key.fingerprint.lower(), KEY_FINGERPRINT.lower())
         self.assertTrue(key.private)
@@ -231,7 +234,7 @@ class KeyManagerKeyManagementTestCase(KeyManagerWithSoledadTestCase):
 
         key = yield self._fetch_key(km, ADDRESS, PUBLIC_KEY)
         self.assertIsInstance(key, OpenPGPKey)
-        self.assertTrue(ADDRESS in key.address)
+        self.assertTrue(ADDRESS in key.uids)
         self.assertEqual(key.validation, ValidationLevels.Provider_Trust)
 
     @defer.inlineCallbacks
@@ -243,7 +246,7 @@ class KeyManagerKeyManagementTestCase(KeyManagerWithSoledadTestCase):
 
         key = yield self._fetch_key(km, ADDRESS_OTHER, PUBLIC_KEY_OTHER)
         self.assertIsInstance(key, OpenPGPKey)
-        self.assertTrue(ADDRESS_OTHER in key.address)
+        self.assertTrue(ADDRESS_OTHER in key.uids)
         self.assertEqual(key.validation, ValidationLevels.Weak_Chain)
 
     def _fetch_key(self, km, address, key):
@@ -273,7 +276,7 @@ class KeyManagerKeyManagementTestCase(KeyManagerWithSoledadTestCase):
         yield km.put_raw_key(PUBLIC_KEY, OpenPGPKey, ADDRESS)
         key = yield km.get_key(ADDRESS, OpenPGPKey)
         self.assertIsInstance(key, OpenPGPKey)
-        self.assertTrue(ADDRESS in key.address)
+        self.assertTrue(ADDRESS in key.uids)
 
     @defer.inlineCallbacks
     def test_fetch_uri_ascii_key(self):
diff --git a/src/leap/keymanager/tests/test_openpgp.py b/src/leap/keymanager/tests/test_openpgp.py
index 8ed049f6..0e5f6bee 100644
--- a/src/leap/keymanager/tests/test_openpgp.py
+++ b/src/leap/keymanager/tests/test_openpgp.py
@@ -335,7 +335,7 @@ class OpenPGPCryptoTestCase(KeyManagerWithSoledadTestCase):
                 deferreds = []
                 for k in (k1, k2, k3, k4, k5):
                     d = self._soledad.create_doc_from_json(
-                        k.get_active_json(ADDRESS))
+                        k.get_active_json())
                     deferreds.append(d)
                 return gatherResults(deferreds)
             elif args[0] == TYPE_FINGERPRINT_PRIVATE_INDEX: