Implement active key document

3 files changed, 161 insertions, 35 deletions

diff --git a/src/leap/keymanager/__init__.py b/src/leap/keymanager/__init__.py
index 0ffb6fc7..562bfbfa 100644
--- a/src/leap/keymanager/__init__.py
+++ b/src/leap/keymanager/__init__.py
@@ -519,15 +519,24 @@ class KeyManager(object):
         except IndexError as e:
             leap_assert(False, "Unsupported key type. Error {0!r}".format(e))
 
-    def put_key(self, key):
+    def put_key(self, key, address=None):
         """
         Put C{key} in local storage.
 
         :param key: The key to be stored
         :type key: EncryptionKey
+        :param address: address for which this key will be active. If not set
+                        all the uids will be activated
+        :type address: str
+
+        :raises KeyAddressMismatch: if address doesn't match any uid on the key
         :raises KeyNotValidUpdate: if a key with the same uid exists and the
                                    new one is not a valid update for it
         """
+        if address is not None and address not in key.address:
+            raise KeyAddressMismatch("UID %s found, but expected %s"
+                                     % (str(key.address), address))
+
         try:
             old_key = self._wrapper_map[type(key)].get_key(key.address[0],
                                                            private=key.private)
@@ -536,7 +545,7 @@ class KeyManager(object):
 
         if key.private or can_upgrade(key, old_key):
             try:
-                self._wrapper_map[type(key)].put_key(key)
+                self._wrapper_map[type(key)].put_key(key, address)
             except IndexError as e:
                 leap_assert(
                     False, "Unsupported key type. Error {0!r}".format(e))
@@ -553,7 +562,7 @@ class KeyManager(object):
         :type key: str
         :param ktype: the type of the key.
         :type ktype: subclass of EncryptionKey
-        :param address: if set used to check that the key is for this address
+        :param address: address for which this key will be active
         :type address: str
         :param validation: validation level for this key
                            (default: 'Weak_Chain')
@@ -564,12 +573,9 @@ class KeyManager(object):
                                    new one is not a valid update for it
         """
         pubkey, _ = self._wrapper_map[ktype].parse_ascii_key(key)
-        if address is not None and address not in pubkey.address:
-            raise KeyAddressMismatch("Key UID %s, but expected %s"
-                                     % (pubkey.address, address))
 
         pubkey.validation = validation
-        self.put_key(pubkey)
+        self.put_key(pubkey, address)
 
     def fetch_key(self, address, uri, ktype,
                   validation=ValidationLevel.Weak_Chain):
@@ -600,12 +606,9 @@ class KeyManager(object):
         pubkey, _ = self._wrapper_map[ktype].parse_ascii_key(res.content)
         if pubkey is None:
             raise KeyNotFound(uri)
-        if address not in pubkey.address:
-            raise KeyAddressMismatch("UID %s found, but expected %s"
-                                     % (str(pubkey.address), address))
 
         pubkey.validation = validation
-        self.put_key(pubkey)
+        self.put_key(pubkey, address)
 
 from ._version import get_versions
 __version__ = get_versions()['version']
diff --git a/src/leap/keymanager/keys.py b/src/leap/keymanager/keys.py
index 5aeb7945..d1b5d933 100644
--- a/src/leap/keymanager/keys.py
+++ b/src/leap/keymanager/keys.py
@@ -64,6 +64,8 @@ KEY_TAGS_KEY = 'tags'
 #
 
 KEYMANAGER_KEY_TAG = 'keymanager-key'
+KEYMANAGER_ACTIVE_TAG = 'keymanager-active'
+KEYMANAGER_ACTIVE_TYPE = '-active'
 
 
 #
@@ -71,12 +73,18 @@ KEYMANAGER_KEY_TAG = 'keymanager-key'
 #
 
 TAGS_PRIVATE_INDEX = 'by-tags-private'
+TYPE_ID_PRIVATE_INDEX = 'by-type-id-private'
 TYPE_ADDRESS_PRIVATE_INDEX = 'by-type-address-private'
 INDEXES = {
     TAGS_PRIVATE_INDEX: [
         KEY_TAGS_KEY,
         'bool(%s)' % KEY_PRIVATE_KEY,
     ],
+    TYPE_ID_PRIVATE_INDEX: [
+        KEY_TYPE_KEY,
+        KEY_ID_KEY,
+        'bool(%s)' % KEY_PRIVATE_KEY,
+    ],
     TYPE_ADDRESS_PRIVATE_INDEX: [
         KEY_TYPE_KEY,
         KEY_ADDRESS_KEY,
@@ -215,6 +223,23 @@ class EncryptionKey(object):
             KEY_TAGS_KEY: [KEYMANAGER_KEY_TAG],
         })
 
+    def get_active_json(self, address):
+        """
+        Return a JSON string describing this key.
+
+        :param address: Address for wich the key is active
+        :type address: str
+        :return: The JSON string describing this key.
+        :rtype: str
+        """
+        return json.dumps({
+            KEY_ADDRESS_KEY: address,
+            KEY_TYPE_KEY: self.__class__.__name__ + KEYMANAGER_ACTIVE_TYPE,
+            KEY_ID_KEY: self.key_id,
+            KEY_PRIVATE_KEY: self.private,
+            KEY_TAGS_KEY: [KEYMANAGER_ACTIVE_TAG],
+        })
+
     def __repr__(self):
         """
         Representation of this class
diff --git a/src/leap/keymanager/openpgp.py b/src/leap/keymanager/openpgp.py
index 38db178c..52655d0e 100644
--- a/src/leap/keymanager/openpgp.py
+++ b/src/leap/keymanager/openpgp.py
@@ -36,9 +36,12 @@ from leap.keymanager.keys import (
     EncryptionScheme,
     is_address,
     build_key_from_dict,
+    TYPE_ID_PRIVATE_INDEX,
     TYPE_ADDRESS_PRIVATE_INDEX,
     KEY_FINGERPRINT_KEY,
     KEY_DATA_KEY,
+    KEY_ID_KEY,
+    KEYMANAGER_ACTIVE_TYPE,
 )
 
 
@@ -193,6 +196,18 @@ def _build_key_from_gpg(address, key, key_data):
     )
 
 
+def _parse_address(address):
+    """
+    Remove the identity suffix after the '+' until the '@'
+    e.g.: test_user+something@provider.com becomes test_user@provider.com
+    since the key belongs to the identity without the '+' suffix.
+
+    :type address: str
+    :rtype: str
+    """
+    return re.sub(r'\+.*\@', '@', address)
+
+
 #
 # The OpenPGP wrapper
 #
@@ -210,7 +225,8 @@ class OpenPGPScheme(EncryptionScheme):
     """
 
     # type used on the soledad documents
-    OPENPGP_KEY_TYPE = OpenPGPKey.__name__
+    KEY_TYPE = OpenPGPKey.__name__
+    ACTIVE_TYPE = KEY_TYPE + KEYMANAGER_ACTIVE_TYPE
 
     def __init__(self, soledad, gpgbinary=None):
         """
@@ -282,7 +298,7 @@ class OpenPGPScheme(EncryptionScheme):
                 openpgp_key = _build_key_from_gpg(
                     address, key,
                     gpg.export_keys(key['fingerprint'], secret=secret))
-                self.put_key(openpgp_key)
+                self.put_key(openpgp_key, address)
 
         return self.get_key(address, private=True)
 
@@ -299,10 +315,7 @@ class OpenPGPScheme(EncryptionScheme):
         :rtype: OpenPGPKey
         @raise KeyNotFound: If the key was not found on local storage.
         """
-        # Remove the identity suffix after the '+' until the '@'
-        # e.g.: test_user+something@provider.com becomes test_user@provider.com
-        # since the key belongs to the identity without the '+' suffix.
-        address = re.sub(r'\+.*\@', '@', address)
+        address = _parse_address(address)
 
         doc = self._get_key_doc(address, private)
         if doc is None:
@@ -371,12 +384,15 @@ class OpenPGPScheme(EncryptionScheme):
 
             return (openpgp_pubkey, openpgp_privkey)
 
-    def put_ascii_key(self, key_data):
+    def put_ascii_key(self, key_data, address=None):
         """
         Put key contained in ascii-armored C{key_data} in local storage.
 
         :param key_data: The key data to be stored.
         :type key_data: str or unicode
+        :param address: address for which this key will be active. If not set
+                        all the uids will be activated
+        :type address: str
         """
         leap_assert_type(key_data, (str, unicode))
 
@@ -387,21 +403,41 @@ class OpenPGPScheme(EncryptionScheme):
             leap_assert(False, repr(e))
 
         if openpgp_pubkey is not None:
-            self.put_key(openpgp_pubkey)
+            self.put_key(openpgp_pubkey, address)
         if openpgp_privkey is not None:
-            self.put_key(openpgp_privkey)
+            self.put_key(openpgp_privkey, address)
 
-    def put_key(self, key):
+    def put_key(self, key, address=None):
         """
         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. If not set
+                        all the uids will be activated
+        :type address: str
         """
-        doc = self._get_key_doc(key.address[0], private=key.private)
-        if doc is None:
-            self._soledad.create_doc_from_json(key.get_json())
+        if address is not None:
+            active_address = [_parse_address(address)]
         else:
+            active_address = key.address
+
+        self._put_key_doc(key)
+        self._put_active_doc(key, active_address)
+
+    def _put_key_doc(self, key):
+        """
+        Put key document in soledad
+
+        :type key: OpenPGPKey
+        """
+        docs = self._soledad.get_from_index(
+            TYPE_ID_PRIVATE_INDEX,
+            self.KEY_TYPE,
+            key.key_id,
+            '1' if key.private else '0')
+        if len(docs) != 0:
+            doc = docs.pop()
             if key.fingerprint == doc.content[KEY_FINGERPRINT_KEY]:
                 # in case of an update of the key merge them with gnupg
                 with self._temporary_gpgwrapper() as gpg:
@@ -412,8 +448,41 @@ class OpenPGPScheme(EncryptionScheme):
                         key.address[0], gpgkey,
                         gpg.export_keys(gpgkey['fingerprint'],
                                         secret=key.private))
-            doc.set_json(key.get_json())
-            self._soledad.put_doc(doc)
+                doc.set_json(key.get_json())
+                self._soledad.put_doc(doc)
+            else:
+                logger.critical(
+                    "Can't put a key whith the same key_id and different "
+                    "fingerprint: %s, %s"
+                    % (key.fingerprint, doc.content[KEY_FINGERPRINT_KEY]))
+        else:
+            self._soledad.create_doc_from_json(key.get_json())
+
+    def _put_active_doc(self, key, addresses):
+        """
+        Put active key document in soledad
+
+        :type key: OpenPGPKey
+        :type addresses: list(str)
+        """
+        for address in addresses:
+            docs = self._soledad.get_from_index(
+                TYPE_ADDRESS_PRIVATE_INDEX,
+                self.ACTIVE_TYPE,
+                address,
+                '1' if key.private else '0')
+            if len(docs) == 1:
+                doc = docs.pop()
+                doc.set_json(key.get_active_json(address))
+                self._soledad.put_doc(doc)
+            else:
+                if len(docs) > 1:
+                    logger.error("There is more than one active key document "
+                                 "for the address %s" % (address,))
+                    for doc in docs:
+                        self._soledad.delete_doc(doc)
+                self._soledad.create_doc_from_json(
+                    key.get_active_json(address))
 
     def _get_key_doc(self, address, private=False):
         """
@@ -428,17 +497,26 @@ class OpenPGPScheme(EncryptionScheme):
         :return: The document with the key or None if it does not exist.
         :rtype: leap.soledad.document.SoledadDocument
         """
-        doclist = self._soledad.get_from_index(
+        activedoc = self._soledad.get_from_index(
             TYPE_ADDRESS_PRIVATE_INDEX,
-            self.OPENPGP_KEY_TYPE,
+            self.ACTIVE_TYPE,
             address,
             '1' if private else '0')
-        if len(doclist) is 0:
+        if len(activedoc) is 0:
             return None
         leap_assert(
+            len(activedoc) is 1,
+            'Found more than one key for address %s!' % (address,))
+
+        key_id = activedoc[0].content[KEY_ID_KEY]
+        doclist = self._soledad.get_from_index(
+            TYPE_ID_PRIVATE_INDEX,
+            self.KEY_TYPE,
+            key_id,
+            '1' if private else '0')
+        leap_assert(
             len(doclist) is 1,
-            'Found more than one %s key for address!' %
-            'private' if private else 'public')
+            'There is %d keys for id %s!' % (len(doclist), key_id))
         return doclist.pop()
 
     def delete_key(self, key):
@@ -447,17 +525,37 @@ class OpenPGPScheme(EncryptionScheme):
 
         May raise:
             errors.KeyNotFound
-            errors.KeyAttributesDiffer
 
         :param key: The key to be removed.
         :type key: EncryptionKey
         """
         leap_assert_type(key, OpenPGPKey)
-        doc = self._get_key_doc(key.address[0], key.private)
+        activedocs = self._soledad.get_from_index(
+            TYPE_ID_PRIVATE_INDEX,
+            self.ACTIVE_TYPE,
+            key.key_id,
+            '1' if key.private else '0')
+        for doc in activedocs:
+            self._soledad.delete_doc(doc)
+
+        docs = self._soledad.get_from_index(
+            TYPE_ID_PRIVATE_INDEX,
+            self.KEY_TYPE,
+            key.key_id,
+            '1' if key.private else '0')
+        if len(docs) == 0:
+            raise errors.KeyNotFound(key)
+        if len(docs) > 1:
+            logger.critical("There is more than one key for key_id %s"
+                            % key.key_id)
+
+        doc = None
+        for d in docs:
+            if d.content['fingerprint'] == key.fingerprint:
+                doc = d
+                break
         if doc is None:
             raise errors.KeyNotFound(key)
-        if doc.content[KEY_FINGERPRINT_KEY] != key.fingerprint:
-            raise errors.KeyAttributesDiffer(key)
         self._soledad.delete_doc(doc)
 
     #