Merge pull request #29 from quantopian/add-crypto-factories

ENH: Add per-user crypto factories.

BUG: Specify notebook<5 and iPy<6 in setup.

Author: nathanwolfe

pgcontents/crypto.py

@@ -1,12 +1,23 @@
 """
 Interface definition for encryption/decryption plugins for
-PostgresContentsManager.
+PostgresContentsManager, and implementations of the interface.
 
 Encryption backends should raise pgcontents.error.CorruptedFile if they
 encounter an input that they cannot decrypt.
 """
+import sys
+import base64
+
+from cryptography.fernet import Fernet
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives import hashes
+from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
+
 from .error import CorruptedFile
 
+if sys.version_info.major == 3:
+    unicode = str
+
 
 class NoEncryption(object):
     """
@@ -127,3 +138,91 @@ def decrypt(self, s):
             except CorruptedFile as e:
                 errors.append(e)
         raise CorruptedFile(errors)
+
+
+def ascii_unicode_to_bytes(v):
+    assert isinstance(v, unicode), "Expected unicode, got %s" % type(v)
+    return v.encode('ascii')
+
+
+def derive_single_fernet_key(password, user_id):
+    """
+    Convert a secret key and a user ID into an encryption key to use with a
+    ``cryptography.fernet.Fernet``.
+
+    Taken from
+    https://cryptography.io/en/latest/fernet/#using-passwords-with-fernet
+
+    Parameters
+    ----------
+    password : unicode
+        ascii-encodable key to derive
+    user_id : unicode
+        ascii-encodable user_id to use as salt
+    """
+    password = ascii_unicode_to_bytes(password)
+    user_id = ascii_unicode_to_bytes(user_id)
+
+    kdf = PBKDF2HMAC(
+        algorithm=hashes.SHA256(),
+        length=32,
+        salt=user_id,
+        iterations=100000,
+        backend=default_backend(),
+    )
+    return base64.urlsafe_b64encode(kdf.derive(password))
+
+
+def derive_fallback_fernet_keys(passwords, user_id):
+    """
+    Derive a list of per-user Fernet keys from a list of master keys and a
+    username.
+
+    If a None is encountered in ``passwords``, it is forwarded.
+
+    Parameters
+    ----------
+    passwords : list[unicode]
+        List of ascii-encodable keys to derive.
+    user_id : unicode or None
+        ascii-encodable user_id to use as salt
+    """
+    # Normally I wouldn't advocate for these kinds of assertions, but we really
+    # really really don't want to mess up deriving encryption keys.
+    assert isinstance(passwords, (list, tuple)), \
+        "Expected list or tuple of keys, got %s." % type(passwords)
+
+    def derive_single_allow_none(k):
+        if k is None:
+            return None
+        return derive_single_fernet_key(k, user_id).decode('ascii')
+
+    return list(map(derive_single_allow_none, passwords))
+
+
+def no_password_crypto_factory():
+    """
+    Create and return a function suitable for passing as a crypto_factory to
+    ``pgcontents.utils.sync.reencrypt_all_users``
+
+    The factory here always returns NoEncryption().  This is useful when passed
+    as ``old_crypto_factory`` to a database that hasn't yet been encrypted.
+    """
+    def factory(user_id):
+        return NoEncryption()
+    return factory
+
+
+def single_password_crypto_factory(password):
+    """
+    Create and return a function suitable for passing as a crypto_factory to
+    ``pgcontents.utils.sync.reencrypt_all_users``
+
+    The factory here returns a ``FernetEncryption`` that uses a key derived
+    from ``password`` and salted with the supplied user_id.
+    """
+    def factory(user_id):
+        return FernetEncryption(
+            Fernet(derive_single_fernet_key(password, user_id))
+        )
+    return factory

pgcontents/tests/test_encryption.py

@@ -0,0 +1,56 @@
+"""
+Tests for notebook encryption utilities.
+"""
+from cryptography.fernet import Fernet
+
+from ..crypto import (
+    derive_fallback_fernet_keys,
+    FallbackCrypto,
+    FernetEncryption,
+    NoEncryption,
+    single_password_crypto_factory,
+)
+
+
+def test_fernet_derivation():
+    pws = [u'currentpassword', u'oldpassword', None]
+
+    # This must be Unicode, so we use the `u` prefix to support py2.
+    user_id = u'4e322fa200fffd0001000001'
+
+    current_crypto = single_password_crypto_factory(pws[0])(user_id)
+    old_crypto = single_password_crypto_factory(pws[1])(user_id)
+
+    def make_single_key_crypto(key):
+        if key is None:
+            return NoEncryption()
+        return FernetEncryption(Fernet(key.encode('ascii')))
+
+    multi_fernet_crypto = FallbackCrypto(
+        [make_single_key_crypto(k)
+         for k in derive_fallback_fernet_keys(pws, user_id)]
+    )
+
+    data = b'ayy lmao'
+
+    # Data encrypted with the current key.
+    encrypted_data_current = current_crypto.encrypt(data)
+    assert encrypted_data_current != data
+    assert current_crypto.decrypt(encrypted_data_current) == data
+
+    # Data encrypted with the old key.
+    encrypted_data_old = old_crypto.encrypt(data)
+    assert encrypted_data_current != data
+    assert old_crypto.decrypt(encrypted_data_old) == data
+
+    # The single fernet with the first key should be able to decrypt the
+    # multi-fernet's encrypted data.
+
+    assert current_crypto.decrypt(multi_fernet_crypto.encrypt(data)) == data
+
+    # Multi should be able decrypt anything encrypted with either key.
+    assert multi_fernet_crypto.decrypt(encrypted_data_current) == data
+    assert multi_fernet_crypto.decrypt(encrypted_data_old) == data
+
+    # Unencrypted data should be returned unchanged.
+    assert multi_fernet_crypto.decrypt(data) == data

pgcontents/utils/ipycompat.py

@@ -46,6 +46,10 @@
         Unicode,
     )
 else:
+    import notebook
+    if notebook.version_info[0] >= 5:
+        raise ImportError("Notebook versions 5 and up are not supported.")
+
     from traitlets.config import Config
     from notebook.services.contents.checkpoints import (
         Checkpoints,

setup.py

@@ -48,7 +48,7 @@ def main():
         extras_require={
             'test': test_reqs,
             'ipy3': ['ipython[test,notebook]<4.0'],
-            'ipy4': ['notebook[test]>=4.0'],
+            'ipy4': ['ipython<6.0', 'notebook[test]>=4.0,<5.0'],
         },
         scripts=[
             'bin/pgcontents',