Add key refreshing for KeyManager.

5 files changed, 218 insertions, 147 deletions

diff --git a/src/leap/common/keymanager/__init__.py b/src/leap/common/keymanager/__init__.py
index f939a4e..82fa99b 100644
--- a/src/leap/common/keymanager/__init__.py
+++ b/src/leap/common/keymanager/__init__.py
@@ -31,6 +31,9 @@ from leap.common.keymanager.errors import (
     KeyNotFound,
     KeyAlreadyExists,
 )
+from leap.common.keymanager.keys import (
+    build_key_from_dict,
+)
 from leap.common.keymanager.openpgp import (
     OpenPGPKey,
     OpenPGPScheme,
@@ -39,6 +42,14 @@ from leap.common.keymanager.openpgp import (
 from leap.common.keymanager.http import HTTPClient
 
 
+TAGS_INDEX = 'by-tags'
+TAGS_AND_PRIVATE_INDEX = 'by-tags-and-private'
+INDEXES = {
+    TAGS_INDEX: ['tags'],
+    TAGS_AND_PRIVATE_INDEX: ['tags', 'bool(private)'],
+}
+
+
 class KeyManager(object):
 
     def __init__(self, address, url, soledad):
@@ -55,10 +66,45 @@ class KeyManager(object):
         """
         self._address = address
         self._http_client = HTTPClient(url)
+        self._soledad = soledad
         self._wrapper_map = {
             OpenPGPKey: OpenPGPScheme(soledad),
             # other types of key will be added to this mapper.
         }
+        self._init_indexes()
+
+    #
+    # utilities
+    #
+
+    def _key_class_from_type(self, ktype):
+        """
+        Return key class from string representation of key type.
+        """
+        return filter(
+            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 send_key(self, ktype, send_private=False, password=None):
         """
@@ -104,7 +150,7 @@ class KeyManager(object):
             json.dumps(data),
             headers)
 
-    def get_key(self, address, ktype, private=False):
+    def get_key(self, address, ktype, private=False, fetch_remote=True):
         """
         Return a key of type C{ktype} bound to C{address}.
 
@@ -129,14 +175,21 @@ class KeyManager(object):
         try:
             return self._wrapper_map[ktype].get_key(address, private=private)
         except KeyNotFound:
-            key = filter(lambda k: isinstance(k, ktype),
-                         self._fetch_keys(address))
-            if key is None:
+            if fetch_remote is False:
+                raise
+            # fetch keys from server and discard unwanted types.
+            keys = filter(lambda k: isinstance(k, ktype),
+                          self.fetch_keys_from_server(address))
+            if len(keys) is 0:
                 raise KeyNotFound()
-            self._wrapper_map[ktype].put_key(key)
+            leap_assert(
+                len(keys) == 1,
+                'Got more than one key of type %s for %s.' %
+                (str(ktype), address))
+            self._wrapper_map[ktype].put_key(keys[0])
             return key
 
-    def _fetch_keys(self, address):
+    def fetch_keys_from_server(self, address):
         """
         Fetch keys bound to C{address} from nickserver.
 
@@ -153,18 +206,42 @@ class KeyManager(object):
         leap_assert(
             keydata['address'] == address,
             "Fetched key for wrong address.")
+        keys = []
         for key in keydata['keys']:
-            # find the key class in the mapper
-            keyCLass = filter(
-                lambda klass: str(klass) == key['type'],
-                self._wrapper_map).pop()
-            yield _build_key_from_dict(kClass, address, key)
+            keys.append(
+                build_key_from_dict(
+                    self._key_class_from_type(key['type']),
+                    address,
+                    key))
+        return keys
+
+    def get_all_keys_in_local_db(self, private=False):
+        """
+        Return all keys stored in local database.
+
+        @return: A list with all keys in local db.
+        @rtype: list
+        """
+        return map(
+            lambda doc: build_key_from_dict(
+                self._key_class_from_type(doc.content['type']),
+                doc.content['address'],
+                doc.content),
+            self._soledad.get_from_index(
+                TAGS_AND_PRIVATE_INDEX,
+                'keymanager-key',
+                '1' if private else '0'))
 
     def refresh_keys(self):
         """
-        Update the user's db of validated keys to see if there are changes.
+        Fetch keys from nickserver and update them locally.
         """
-        raise NotImplementedError(self.refresh_keys)
+        addresses = set(map(
+            lambda doc: doc.address,
+            self.get_all_keys_in_local_db(False)))
+        for address in addresses:
+            for key in self.fetch_keys_from_server(address):
+                self._wrapper_map[key.__class__].put_key(key)
 
     def gen_key(self, ktype):
         """
diff --git a/src/leap/common/keymanager/keys.py b/src/leap/common/keymanager/keys.py
index 250c2fa..453e0ed 100644
--- a/src/leap/common/keymanager/keys.py
+++ b/src/leap/common/keymanager/keys.py
@@ -25,11 +25,81 @@ try:
     import simplejson as json
 except ImportError:
     import json  # noqa
+import re
 
 
+from hashlib import sha256
 from abc import ABCMeta, abstractmethod
+from leap.common.check import leap_assert
 
 
+#
+# Key handling utilities
+#
+
+def is_address(address):
+    """
+    Return whether the given C{address} is in the form user@provider.
+
+    @param address: The address to be tested.
+    @type address: str
+    @return: Whether C{address} is in the form user@provider.
+    @rtype: bool
+    """
+    return bool(re.match('[\w.-]+@[\w.-]+', address))
+
+
+def build_key_from_dict(kClass, address, kdict):
+    """
+    Build an C{kClass} key bound to C{address} based on info in C{kdict}.
+
+    @param address: The address bound to the key.
+    @type address: str
+    @param kdict: Dictionary with key data.
+    @type kdict: dict
+    @return: An instance of the key.
+    @rtype: C{kClass}
+    """
+    leap_assert(address == kdict['address'], 'Wrong address in key data.')
+    return kClass(
+        address,
+        key_id=kdict['key_id'],
+        fingerprint=kdict['fingerprint'],
+        key_data=kdict['key_data'],
+        private=kdict['private'],
+        length=kdict['length'],
+        expiry_date=kdict['expiry_date'],
+        first_seen_at=kdict['first_seen_at'],
+        last_audited_at=kdict['last_audited_at'],
+        validation=kdict['validation'],  # TODO: verify for validation.
+    )
+
+
+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 = 'private' if private else 'public'
+    return sha256('keymanager-'+address+'-'+ktype+'-'+visibility).hexdigest()
+
+
+#
+# Abstraction for encryption keys
+#
+
 class EncryptionKey(object):
     """
     Abstract class for encryption keys.
@@ -82,6 +152,7 @@ class EncryptionKey(object):
             'validation': self.validation,
             'first_seen_at': self.first_seen_at,
             'last_audited_at': self.last_audited_at,
+            'tags': ['keymanager-key'],
         })
 
 
diff --git a/src/leap/common/keymanager/openpgp.py b/src/leap/common/keymanager/openpgp.py
index ace8c1e..fa3f732 100644
--- a/src/leap/common/keymanager/openpgp.py
+++ b/src/leap/common/keymanager/openpgp.py
@@ -33,13 +33,11 @@ from leap.common.keymanager.errors import (
 from leap.common.keymanager.keys import (
     EncryptionKey,
     EncryptionScheme,
+    is_address,
+    keymanager_doc_id,
+    build_key_from_dict,
 )
 from leap.common.keymanager.gpg import GPGWrapper
-from leap.common.keymanager.util import (
-    _is_address,
-    _build_key_from_doc,
-    _keymanager_doc_id,
-)
 
 
 #
@@ -307,7 +305,7 @@ class OpenPGPScheme(EncryptionScheme):
         @raise KeyAlreadyExists: If key already exists in local database.
         """
         # make sure the key does not already exist
-        leap_assert(_is_address(address), 'Not an user address: %s' % address)
+        leap_assert(is_address(address), 'Not an user address: %s' % address)
         try:
             self.get_key(address)
             raise KeyAlreadyExists(address)
@@ -358,11 +356,11 @@ class OpenPGPScheme(EncryptionScheme):
         @rtype: OpenPGPKey
         @raise KeyNotFound: If the key was not found on local storage.
         """
-        leap_assert(_is_address(address), 'Not an user address: %s' % address)
+        leap_assert(is_address(address), 'Not an user address: %s' % address)
         doc = self._get_key_doc(address, private)
         if doc is None:
             raise KeyNotFound(address)
-        return _build_key_from_doc(OpenPGPKey, address, doc)
+        return build_key_from_dict(OpenPGPKey, address, doc.content)
 
     def put_key_raw(self, data):
         """
@@ -422,7 +420,7 @@ class OpenPGPScheme(EncryptionScheme):
         if doc is None:
             self._soledad.create_doc_from_json(
                 key.get_json(),
-                doc_id=_keymanager_doc_id(
+                doc_id=keymanager_doc_id(
                     OpenPGPKey, key.address, key.private))
         else:
             doc.set_json(key.get_json())
@@ -442,7 +440,7 @@ class OpenPGPScheme(EncryptionScheme):
         @rtype: leap.soledad.backends.leap_backend.LeapDocument
         """
         return self._soledad.get_doc(
-            _keymanager_doc_id(OpenPGPKey, address, private))
+            keymanager_doc_id(OpenPGPKey, address, private))
 
     def delete_key(self, key):
         """
@@ -458,5 +456,5 @@ class OpenPGPScheme(EncryptionScheme):
         if stored_key.__dict__ != key.__dict__:
             raise KeyAttributesDiffer(key)
         doc = self._soledad.get_doc(
-            _keymanager_doc_id(OpenPGPKey, key.address, key.private))
+            keymanager_doc_id(OpenPGPKey, key.address, key.private))
         self._soledad.delete_doc(doc)
diff --git a/src/leap/common/keymanager/util.py b/src/leap/common/keymanager/util.py
deleted file mode 100644
index 667d2b2..0000000
--- a/src/leap/common/keymanager/util.py
+++ /dev/null
@@ -1,103 +0,0 @@
-# -*- coding: utf-8 -*-
-# util.py
-# Copyright (C) 2013 LEAP
-#
-# This program is free software: you can redistribute it and/or modify
-# it under the terms of the GNU General Public License as published by
-# the Free Software Foundation, either version 3 of the License, or
-# (at your option) any later version.
-#
-# This program is distributed in the hope that it will be useful,
-# but WITHOUT ANY WARRANTY; without even the implied warranty of
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
-# GNU General Public License for more details.
-#
-# You should have received a copy of the GNU General Public License
-# along with this program. If not, see <http://www.gnu.org/licenses/>.
-
-
-"""
-Utilities for the Key Manager.
-"""
-
-
-import re
-
-
-from hashlib import sha256
-from binascii import b2a_base64
-
-
-from leap.common.check import leap_assert
-
-
-def _is_address(address):
-    """
-    Return whether the given C{address} is in the form user@provider.
-
-    @param address: The address to be tested.
-    @type address: str
-    @return: Whether C{address} is in the form user@provider.
-    @rtype: bool
-    """
-    return bool(re.match('[\w.-]+@[\w.-]+', address))
-
-
-def _build_key_from_dict(kClass, address, kdict):
-    """
-    Build an C{kClass} key bound to C{address} based on info in C{kdict}.
-
-    @param address: The address bound to the key.
-    @type address: str
-    @param kdict: Dictionary with key data.
-    @type kdict: dict
-    @return: An instance of the key.
-    @rtype: C{kClass}
-    """
-    return kClass(
-        address,
-        key_id=kdict['key_id'],
-        fingerprint=kdict['fingerprint'],
-        key_data=kdict['key_data'],
-        private=kdict['private'],
-        length=kdict['length'],
-        expiry_date=kdict['expiry_date'],
-        first_seen_at=kdict['first_seen_at'],
-        last_audited_at=kdict['last_audited_at'],
-        validation=kdict['validation'],  # TODO: verify for validation.
-    )
-
-
-def _build_key_from_doc(kClass, address, doc):
-    """
-    Build an C{kClass} for C{address} based on C{doc} from local storage.
-
-    @param address: The address bound to the key.
-    @type address: str
-    @param doc: Document obtained from Soledad storage.
-    @type doc: leap.soledad.backends.leap_backend.LeapDocument
-    @return: An instance of the key.
-    @rtype: C{kClass}
-    """
-    return _build_key_from_dict(kClass, address, doc.content)
-
-
-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 = 'private' if private else 'public'
-    return sha256('key-manager-'+address+'-'+ktype+'-'+visibility).hexdigest()
diff --git a/src/leap/common/tests/test_keymanager.py b/src/leap/common/tests/test_keymanager.py
index f9b478f..9bafb89 100644
--- a/src/leap/common/tests/test_keymanager.py
+++ b/src/leap/common/tests/test_keymanager.py
@@ -29,12 +29,18 @@ from leap.soledad import Soledad
 from leap.soledad.crypto import SoledadCrypto
 
 
-from leap.common.keymanager import KeyManager, openpgp, KeyNotFound
+from leap.common.keymanager import (
+    KeyManager,
+    openpgp,
+    KeyNotFound,
+    TAGS_INDEX,
+    TAGS_AND_PRIVATE_INDEX,
+)
 from leap.common.keymanager.openpgp import OpenPGPKey
-from leap.common.keymanager.util import (
-    _is_address,
-    _build_key_from_dict,
-    _keymanager_doc_id,
+from leap.common.keymanager.keys import (
+    is_address,
+    build_key_from_dict,
+    keymanager_doc_id,
 )
 from leap.common.keymanager import errors
 
@@ -47,21 +53,21 @@ class KeyManagerUtilTestCase(BaseLeapTest):
     def tearDown(self):
         pass
 
-    def test__is_address(self):
+    def test_is_address(self):
         self.assertTrue(
-            _is_address('user@leap.se'),
+            is_address('user@leap.se'),
             'Incorrect address detection.')
         self.assertFalse(
-            _is_address('userleap.se'),
+            is_address('userleap.se'),
             'Incorrect address detection.')
         self.assertFalse(
-            _is_address('user@'),
+            is_address('user@'),
             'Incorrect address detection.')
         self.assertFalse(
-            _is_address('@leap.se'),
+            is_address('@leap.se'),
             'Incorrect address detection.')
 
-    def test__build_key_from_dict(self):
+    def test_build_key_from_dict(self):
         kdict = {
             'address': 'leap@leap.se',
             'key_id': 'key_id',
@@ -74,7 +80,7 @@ class KeyManagerUtilTestCase(BaseLeapTest):
             'last_audited_at': 'last_audited_at',
             'validation': 'validation',
         }
-        key = _build_key_from_dict(OpenPGPKey, 'leap@leap.se', kdict)
+        key = build_key_from_dict(OpenPGPKey, 'leap@leap.se', kdict)
         self.assertEqual(
             kdict['address'], key.address,
             'Wrong data in key.')
@@ -106,14 +112,14 @@ class KeyManagerUtilTestCase(BaseLeapTest):
             kdict['validation'], key.validation,
             'Wrong data in key.')
 
-    def test__keymanager_doc_id(self):
-        doc_id1 = _keymanager_doc_id(
+    def test_keymanager_doc_id(self):
+        doc_id1 = keymanager_doc_id(
             OpenPGPKey, 'leap@leap.se', private=False)
-        doc_id2 = _keymanager_doc_id(
+        doc_id2 = keymanager_doc_id(
             OpenPGPKey, 'leap@leap.se', private=True)
-        doc_id3 = _keymanager_doc_id(
+        doc_id3 = keymanager_doc_id(
             OpenPGPKey, 'user@leap.se', private=False)
-        doc_id4 = _keymanager_doc_id(
+        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!')
@@ -123,7 +129,7 @@ class KeyManagerUtilTestCase(BaseLeapTest):
         self.assertFalse(doc_id3 == doc_id4, 'Doc ids are equal!')
 
 
-class KeyManagerCryptoTestCase(BaseLeapTest):
+class KeyManagerWithSoledadTestCase(BaseLeapTest):
 
     def setUp(self):
         self._soledad = Soledad(
@@ -144,8 +150,9 @@ class KeyManagerCryptoTestCase(BaseLeapTest):
     def tearDown(self):
         pass
 
-    def _key_manager(user='user@leap.se', url='https://domain.org:6425'):
-        return KeyManager(user, url)
+
+
+class OpenPGPCryptoTestCase(KeyManagerWithSoledadTestCase):
 
     def _test_openpgp_gen_key(self):
         pgp = openpgp.OpenPGPScheme(self._soledad)
@@ -229,6 +236,27 @@ class KeyManagerCryptoTestCase(BaseLeapTest):
         self.assertEqual('data', plaintext)
 
 
+class KeyManagerKeyManagementTestCase(
+    KeyManagerWithSoledadTestCase):
+
+    def _key_manager(self, user='leap@leap.se', url=''):
+        return KeyManager(user, url, self._soledad)
+
+    def test_get_all_keys_in_db(self):
+        km = self._key_manager()
+        km._wrapper_map[OpenPGPKey].put_key_raw(PRIVATE_KEY)
+        # get public keys
+        keys = km.get_all_keys_in_local_db(False)
+        self.assertEqual(len(keys), 1, 'Wrong number of keys')
+        self.assertEqual('leap@leap.se', keys[0].address)
+        self.assertFalse(keys[0].private)
+        # get private keys
+        keys = km.get_all_keys_in_local_db(True)
+        self.assertEqual(len(keys), 1, 'Wrong number of keys')
+        self.assertEqual('leap@leap.se', keys[0].address)
+        self.assertTrue(keys[0].private)
+
+
 # Key material for testing
 KEY_FINGERPRINT = "E36E738D69173C13D709E44F2F455E2824D18DDF"
 PUBLIC_KEY = """