Use indexes to fetch keys.

4 files changed, 59 insertions, 84 deletions

diff --git a/src/leap/common/keymanager/__init__.py b/src/leap/common/keymanager/__init__.py
index 7aaeddf..b6bef21 100644
--- a/src/leap/common/keymanager/__init__.py
+++ b/src/leap/common/keymanager/__init__.py
@@ -35,6 +35,7 @@ from leap.common.keymanager.errors import (
 from leap.common.keymanager.keys import (
     build_key_from_dict,
     KEYMANAGER_KEY_TAG,
+    TAGS_PRIVATE_INDEX,
 )
 from leap.common.keymanager.openpgp import (
     OpenPGPKey,
@@ -43,18 +44,6 @@ from leap.common.keymanager.openpgp import (
 
 
 #
-# key indexing constants.
-#
-
-TAGS_INDEX = 'by-tags'
-TAGS_AND_PRIVATE_INDEX = 'by-tags-and-private'
-INDEXES = {
-    TAGS_INDEX: ['tags'],
-    TAGS_AND_PRIVATE_INDEX: ['tags', 'bool(private)'],
-}
-
-
-#
 # The Key Manager
 #
 
@@ -103,8 +92,6 @@ class KeyManager(object):
             OpenPGPKey: OpenPGPScheme(soledad),
             # other types of key will be added to this mapper.
         }
-        # initialize the indexes needed to query the database
-        self._init_indexes()
         # the following are used to perform https requests
         self._fetcher = requests
         self._session = self._fetcher.session()
@@ -121,26 +108,6 @@ class KeyManager(object):
             lambda klass: str(klass) == ktype,
             self._wrapper_map).pop()
 
-    def _init_indexes(self):
-        """
-        Initialize the database indexes.
-        """
-        # Ask the database for currently existing indexes.
-        db_indexes = dict(self._soledad.list_indexes())
-        # Loop through the indexes we expect to find.
-        for name, expression in INDEXES.items():
-            if name not in db_indexes:
-                # The index does not yet exist.
-                self._soledad.create_index(name, *expression)
-                continue
-            if expression == db_indexes[name]:
-                # The index exists and is up to date.
-                continue
-            # The index exists but the definition is not what expected, so we
-            # delete it and add the proper index expression.
-            self._soledad.delete_index(name)
-            self._soledad.create_index(name, *expression)
-
     def _get(self, uri, data=None):
         """
         Send a GET request to C{uri} containing C{data}.
@@ -294,7 +261,7 @@ class KeyManager(object):
                 doc.content['address'],
                 doc.content),
             self._soledad.get_from_index(
-                TAGS_AND_PRIVATE_INDEX,
+                TAGS_PRIVATE_INDEX,
                 KEYMANAGER_KEY_TAG,
                 '1' if private else '0'))
 
diff --git a/src/leap/common/keymanager/keys.py b/src/leap/common/keymanager/keys.py
index 1d87858..be1f113 100644
--- a/src/leap/common/keymanager/keys.py
+++ b/src/leap/common/keymanager/keys.py
@@ -28,7 +28,6 @@ except ImportError:
 import re
 
 
-from hashlib import sha256
 from abc import ABCMeta, abstractmethod
 from leap.common.check import leap_assert
 
@@ -57,6 +56,26 @@ KEY_TAGS_KEY = 'tags'
 
 KEYMANAGER_KEY_TAG = 'keymanager-key'
 
+
+#
+# key indexing constants.
+#
+
+TAGS_PRIVATE_INDEX = 'by-tags-private'
+TAGS_ADDRESS_PRIVATE_INDEX = 'by-tags-address-private'
+INDEXES = {
+    TAGS_PRIVATE_INDEX: [
+        KEY_TAGS_KEY,
+        'bool(%s)' % KEY_PRIVATE_KEY,
+    ],
+    TAGS_ADDRESS_PRIVATE_INDEX: [
+        KEY_TAGS_KEY,
+        KEY_ADDRESS_KEY,
+        'bool(%s)' % KEY_PRIVATE_KEY,
+    ]
+}
+
+
 #
 # Key handling utilities
 #
@@ -101,27 +120,6 @@ def build_key_from_dict(kClass, address, kdict):
     )
 
 
-def keymanager_doc_id(ktype, address, private=False):
-    """
-    Return the document id for the document containing a key for
-    C{address}.
-
-    @param address: The type of the key.
-    @type address: KeyType
-    @param address: The address bound to the key.
-    @type address: str
-    @param private: Whether the key is private or not.
-    @type private: bool
-    @return: The document id for the document that stores a key bound to
-        C{address}.
-    @rtype: str
-    """
-    leap_assert(is_address(address), "Wrong address format: %s" % address)
-    ktype = str(ktype)
-    visibility = KEY_PRIVATE_KEY if private else 'public'
-    return sha256('keymanager-'+address+'-'+ktype+'-'+visibility).hexdigest()
-
-
 #
 # Abstraction for encryption keys
 #
@@ -215,6 +213,27 @@ class EncryptionScheme(object):
         @type soledad: leap.soledad.Soledad
         """
         self._soledad = soledad
+        self._init_indexes()
+
+    def _init_indexes(self):
+        """
+        Initialize the database indexes.
+        """
+        # Ask the database for currently existing indexes.
+        db_indexes = dict(self._soledad.list_indexes())
+        # Loop through the indexes we expect to find.
+        for name, expression in INDEXES.items():
+            if name not in db_indexes:
+                # The index does not yet exist.
+                self._soledad.create_index(name, *expression)
+                continue
+            if expression == db_indexes[name]:
+                # The index exists and is up to date.
+                continue
+            # The index exists but the definition is not what expected, so we
+            # delete it and add the proper index expression.
+            self._soledad.delete_index(name)
+            self._soledad.create_index(name, *expression)
 
     @abstractmethod
     def get_key(self, address, private=False):
diff --git a/src/leap/common/keymanager/openpgp.py b/src/leap/common/keymanager/openpgp.py
index d53afd6..db9efeb 100644
--- a/src/leap/common/keymanager/openpgp.py
+++ b/src/leap/common/keymanager/openpgp.py
@@ -26,13 +26,13 @@ import tempfile
 
 from leap.common.check import leap_assert, leap_assert_type
 from leap.common.keymanager import errors
-
 from leap.common.keymanager.keys import (
     EncryptionKey,
     EncryptionScheme,
     is_address,
-    keymanager_doc_id,
     build_key_from_dict,
+    KEYMANAGER_KEY_TAG,
+    TAGS_ADDRESS_PRIVATE_INDEX,
 )
 from leap.common.keymanager.gpg import GPGWrapper
 
@@ -588,10 +588,7 @@ class OpenPGPScheme(EncryptionScheme):
         """
         doc = self._get_key_doc(key.address, private=key.private)
         if doc is None:
-            self._soledad.create_doc_from_json(
-                key.get_json(),
-                doc_id=keymanager_doc_id(
-                    OpenPGPKey, key.address, key.private))
+            self._soledad.create_doc_from_json(key.get_json())
         else:
             doc.set_json(key.get_json())
             self._soledad.put_doc(doc)
@@ -609,8 +606,18 @@ class OpenPGPScheme(EncryptionScheme):
         @return: The document with the key or None if it does not exist.
         @rtype: leap.soledad.backends.leap_backend.LeapDocument
         """
-        return self._soledad.get_doc(
-            keymanager_doc_id(OpenPGPKey, address, private))
+        doclist = self._soledad.get_from_index(
+            TAGS_ADDRESS_PRIVATE_INDEX,
+            KEYMANAGER_KEY_TAG,
+            address,
+            '1' if private else '0')
+        if len(doclist) is 0:
+            return None
+        leap_assert(
+            len(doclist) is 1,
+            'Found more than one %s key for address!' %
+            'private' if private else 'public')
+        return doclist.pop()
 
     def delete_key(self, key):
         """
@@ -625,6 +632,5 @@ class OpenPGPScheme(EncryptionScheme):
             raise errors.KeyNotFound(key)
         if stored_key.__dict__ != key.__dict__:
             raise errors.KeyAttributesDiffer(key)
-        doc = self._soledad.get_doc(
-            keymanager_doc_id(OpenPGPKey, key.address, key.private))
+        doc = self._get_key_doc(key.address, key.private)
         self._soledad.delete_doc(doc)
diff --git a/src/leap/common/tests/test_keymanager.py b/src/leap/common/tests/test_keymanager.py
index cffa073..73611b6 100644
--- a/src/leap/common/tests/test_keymanager.py
+++ b/src/leap/common/tests/test_keymanager.py
@@ -44,7 +44,6 @@ from leap.common.keymanager.openpgp import OpenPGPKey
 from leap.common.keymanager.keys import (
     is_address,
     build_key_from_dict,
-    keymanager_doc_id,
 )
 from leap.common.keymanager import errors
 
@@ -120,22 +119,6 @@ class KeyManagerUtilTestCase(BaseLeapTest):
             kdict['validation'], key.validation,
             'Wrong data in key.')
 
-    def test_keymanager_doc_id(self):
-        doc_id1 = keymanager_doc_id(
-            OpenPGPKey, ADDRESS, private=False)
-        doc_id2 = keymanager_doc_id(
-            OpenPGPKey, ADDRESS, private=True)
-        doc_id3 = keymanager_doc_id(
-            OpenPGPKey, 'user@leap.se', private=False)
-        doc_id4 = keymanager_doc_id(
-            OpenPGPKey, 'user@leap.se', private=True)
-        self.assertFalse(doc_id1 == doc_id2, 'Doc ids are equal!')
-        self.assertFalse(doc_id1 == doc_id3, 'Doc ids are equal!')
-        self.assertFalse(doc_id1 == doc_id4, 'Doc ids are equal!')
-        self.assertFalse(doc_id2 == doc_id3, 'Doc ids are equal!')
-        self.assertFalse(doc_id2 == doc_id4, 'Doc ids are equal!')
-        self.assertFalse(doc_id3 == doc_id4, 'Doc ids are equal!')
-
 
 class KeyManagerWithSoledadTestCase(BaseLeapTest):