Allow passing Fernet to Encrypted Cookie Storage (#448)

Author: Dreamsorcerer

README.rst

@@ -32,7 +32,6 @@ A trivial usage example:
 .. code:: python
 
     import time
-    import base64
     from cryptography import fernet
     from aiohttp import web
     from aiohttp_session import setup, get_session
@@ -49,10 +48,9 @@ A trivial usage example:
 
     def make_app():
         app = web.Application()
-        # secret_key must be 32 url-safe base64-encoded bytes
         fernet_key = fernet.Fernet.generate_key()
-        secret_key = base64.urlsafe_b64decode(fernet_key)
-        setup(app, EncryptedCookieStorage(secret_key))
+        fernet = fernet.Fernet(fernet_key)
+        setup(app, EncryptedCookieStorage(fernet))
         app.router.add_get('/', handler)
         return app
 

aiohttp_session/cookie_storage.py

@@ -15,7 +15,7 @@ class EncryptedCookieStorage(AbstractStorage):
 
     def __init__(
         self,
-        secret_key: Union[str, bytes, bytearray],
+        secret_key: Union[str, bytes, bytearray, fernet.Fernet],
         *,
         cookie_name: str = "AIOHTTP_SESSION",
         domain: Optional[str] = None,
@@ -39,11 +39,12 @@ def __init__(
             decoder=decoder,
         )
 
-        if isinstance(secret_key, str):
-            pass
-        elif isinstance(secret_key, (bytes, bytearray)):
-            secret_key = base64.urlsafe_b64encode(secret_key)
-        self._fernet = fernet.Fernet(secret_key)
+        if isinstance(secret_key, fernet.Fernet):
+            self._fernet = secret_key
+        else:
+            if isinstance(secret_key, (bytes, bytearray)):
+                secret_key = base64.urlsafe_b64encode(secret_key)
+            self._fernet = fernet.Fernet(secret_key)
 
     async def load_session(self, request: web.Request) -> Session:
         cookie = self.load_cookie(request)

tests/test_encrypted_cookie_storage.py

@@ -32,7 +32,7 @@ def make_cookie(client: TestClient, fernet: Fernet, data: Dict[str, Any]) -> Non
     client.session.cookie_jar.update_cookies({"AIOHTTP_SESSION": encrypted_data})
 
 
-def create_app(handler: Handler, key: Union[str, bytes, bytearray]) -> web.Application:
+def create_app(handler: Handler, key: Union[str, bytes, bytearray, Fernet]) -> web.Application:
     middleware = session_middleware(EncryptedCookieStorage(key))
     app = web.Application(middlewares=[middleware])
     app.router.add_route("GET", "/", handler)
@@ -102,6 +102,23 @@ async def handler(request: web.Request) -> web.StreamResponse:
     assert resp.status == 200
 
 
+async def test_load_existing_session_with_fernet(
+  aiohttp_client: AiohttpClient, fernet: Fernet
+) -> None:
+    async def handler(request: web.Request) -> web.StreamResponse:
+        session = await get_session(request)
+        assert isinstance(session, Session)
+        assert not session.new
+        assert not session._changed
+        assert {"a": 1, "b": 12} == session  # type: ignore[comparison-overlap]
+        return web.Response(body=b"OK")
+
+    client = await aiohttp_client(create_app(handler, fernet))
+    make_cookie(client, fernet, {"a": 1, "b": 12})
+    resp = await client.get("/")
+    assert resp.status == 200
+
+
 async def test_change_session(
     aiohttp_client: AiohttpClient, fernet: Fernet, key: bytes
 ) -> None: