Merge pull request #16 from graingert/cryptography

use simpler cryptography library

.travis.yml

@@ -11,7 +11,7 @@ install:
 - pip install coveralls --use-mirrors
 - pip install aiohttp
 - pip install aioredis
-- pip install pycrypto
+- pip install cryptography
 - python setup.py develop
 
 script:

README.rst

@@ -62,9 +62,9 @@ Available session storages are:
   encodes it via AES cipher. ``secrect_key`` is a ``bytes`` key for AES
   encryption/decryption, the length should be 16 bytes.
 
-  Requires ``PyCrypto`` library::
+  Requires ``crypotgraphy`` library::
 
-      $ pip install aiohttp_session[pycrypto]
+      $ pip install aiohttp_session[secure]
 
 * ``aiohttp_session.redis_storage.RedisStorage(redis_pool)`` -- stores
   JSON-ed data into *redis*, keepeng into cookie only redis key

aiohttp_session/cookie_storage.py

@@ -1,10 +1,10 @@
 import asyncio
 import json
 import base64
-from . import AbstractStorage, Session
 
-from Crypto.Cipher import AES
-from Crypto import Random
+from cryptography import fernet
+
+from . import AbstractStorage, Session
 
 
 class EncryptedCookieStorage(AbstractStorage):
@@ -18,41 +18,28 @@ def __init__(self, secret_key, *, cookie_name="AIOHTTP_SESSION",
                          max_age=max_age, path=path, secure=secure,
                          httponly=httponly)
 
-        self._secret_key = secret_key
-        if len(self._secret_key) % AES.block_size != 0:
-            raise TypeError(
-                'Secret key must be a multiple of {} in length'.format(
-                    AES.block_size))
+        self._fernet = fernet.Fernet(base64.urlsafe_b64encode(secret_key))
 
     @asyncio.coroutine
     def load_session(self, request):
         cookie = self.load_cookie(request)
         if cookie is None:
             return Session(None, new=True)
         else:
-            cookie = base64.b64decode(cookie)
-            iv = cookie[:AES.block_size]
-            data = cookie[AES.block_size:]
-            cipher = AES.new(self._secret_key, AES.MODE_CBC, iv)
-            decrypted = cipher.decrypt(data)
-            data = json.loads(decrypted.decode('utf-8'))
+            data = json.loads(
+                self._fernet.decrypt(cookie.encode('utf-8')).decode('utf-8')
+            )
             return Session(None, data=data, new=False)
 
     @asyncio.coroutine
     def save_session(self, request, response, session):
         if session.empty:
             return self.save_cookie(response, session._mapping)
 
-        cookie_data = json.dumps(self._get_session_data(session)).encode(
-            'utf-8')
-        if len(cookie_data) % AES.block_size != 0:
-            # padding with spaces to full blocks
-            to_pad = AES.block_size - (len(cookie_data) % AES.block_size)
-            cookie_data += b' ' * to_pad
-
-        iv = Random.new().read(AES.block_size)
-        cipher = AES.new(self._secret_key, AES.MODE_CBC, iv)
-        encrypted = cipher.encrypt(cookie_data)
-        encrypted = iv + encrypted
-        b64coded = base64.b64encode(encrypted).decode('utf-8')
-        self.save_cookie(response, b64coded)
+        cookie_data = json.dumps(
+            self._get_session_data(session)
+        ).encode('utf-8')
+        self.save_cookie(
+            response,
+            self._fernet.encrypt(cookie_data).decode('utf-8'),
+        )

docs/index.rst

@@ -63,9 +63,9 @@ Available session storages are:
   encodes it via AES cipher. ``secrect_key`` is a ``bytes`` key for AES
   encryption/decryption, the length should be 16 bytes.
 
-  Requires ``PyCrypto`` library::
+  Requires ``cryptography`` library::
 
-      $ pip install aiohttp_session[pycrypto]
+      $ pip install aiohttp_session[secure]
 
 * ``aiohttp_session.redis_storage.RedisStorage(redis_pool)`` -- stores
   JSON-ed data into *redis*, keepeng into cookie only redis key

setup.py

@@ -17,8 +17,11 @@ def read(f):
 
 install_requires = ['aiohttp>=0.14']
 tests_require = install_requires + ['nose']
-extras_require = {'aioredis': ['aioredis>=0.1.4'],
-                  'pycrypto': ['pycrypto']}
+extras_require = {
+    'aioredis': ['aioredis>=0.1.4'],
+    'pycrypto': ['cryptography'],
+    'secure': ['cryptography'],
+}
 
 
 setup(name='aiohttp_session',

tests/test_encrypted_cookie_storage.py

@@ -5,10 +5,10 @@
 import base64
 import time
 
-from Crypto.Cipher import AES
-from Crypto import Random
-
 from aiohttp import web, request
+
+from cryptography import fernet
+
 from aiohttp_session import Session, session_middleware, get_session
 from aiohttp_session.cookie_storage import EncryptedCookieStorage
 
@@ -18,7 +18,10 @@ class TestEncryptedCookieStorage(unittest.TestCase):
     def setUp(self):
         self.loop = asyncio.new_event_loop()
         asyncio.set_event_loop(None)
-        self.key = b'Sixteen byte key'
+
+        key = fernet.Fernet.generate_key()
+        self.fernet = fernet.Fernet(key)
+        self.key = base64.urlsafe_b64decode(key)
         self.handler = None
         self.srv = None
 
@@ -63,28 +66,18 @@ def make_cookie(self, data):
             session_data = data
 
         cookie_data = json.dumps(session_data).encode('utf-8')
-        if len(cookie_data) % AES.block_size != 0:
-            # padding with spaces to full blocks
-            to_pad = AES.block_size - (len(cookie_data) % AES.block_size)
-            cookie_data += b' ' * to_pad
-        iv = Random.new().read(AES.block_size)
-        cipher = AES.new(self.key, AES.MODE_CBC, iv)
-        encrypted = cipher.encrypt(cookie_data)
-        encrypted = iv + encrypted
-        b64coded = base64.b64encode(encrypted).decode('utf-8')
-        return {'AIOHTTP_SESSION': b64coded}
+        data = self.fernet.encrypt(cookie_data).decode('utf-8')
+
+        return {'AIOHTTP_SESSION': data}
 
     def decrypt(self, cookie_value):
         assert type(cookie_value) == str
-        decoded = base64.b64decode(cookie_value)
-        iv = decoded[:AES.block_size]
-        data = decoded[AES.block_size:]
-        cipher = AES.new(self.key, AES.MODE_CBC, iv)
-        decrypted = cipher.decrypt(data).decode('utf-8')
-        return json.loads(decrypted)
+        return json.loads(
+            self.fernet.decrypt(cookie_value.encode('utf-8')).decode('utf-8')
+        )
 
     def test_invalid_key(self):
-        with self.assertRaises(TypeError):
+        with self.assertRaises(ValueError):
             EncryptedCookieStorage(b'123')  # short key
 
     def test_create_new_sesssion(self):