APIStar Cookie Sessions

example.py

import contextlib
import typing
import json
import os

from werkzeug.http import dump_cookie, parse_cookie

from apistar import http
from apistar.interfaces import SessionStore

from cryptography.fernet import Fernet

COOKIE_NAME = 'session'

default_cipher_key = 'hdJOXy9nenVk3vllaxtbM7WyUL3UaykReMgrkCnmyxc='
cipher = Fernet(os.getenv('SESSION_CIPHER_KEY', default_cipher_key))


class CookieSessionStore(SessionStore):

    def new(self) -> http.Session:
        return http.Session(session_id='')

    def load(self, session_data: str) -> http.Session:
        try:
            data = json.loads(cipher.decrypt(str.encode(session_data)))
        except (ValueError, Exception):
            return self.new()
        return http.Session(session_id='', data=data)

    def save(self, session: http.Session):
        data = cipher.encrypt(str.encode(json.dumps(session.data)))
        cookie = dump_cookie(COOKIE_NAME, data)
        return {'Set-Cookie': cookie}


@contextlib.contextmanager
def init_cookie_session(cookie: http.Header,
                        response_headers: http.ResponseHeaders) -> typing.Generator[http.Session, None, None]:

    cookies = parse_cookie(cookie)
    session_data = cookies.get(COOKIE_NAME)
    if session_data:
        session = CookieSessionStore().load(session_data)
    else:
        session = CookieSessionStore().new()
    try:
        yield session
    finally:
        session_headers = CookieSessionStore().save(session)
        response_headers.update(session_headers)


        
        
# Other file        
from utils.cookie_sessions import init_cookie_session


# Cookie based session
cookie_sessions = [
    Component(http.Session, init=init_cookie_session)
]


components = cookie_sessions


# I also have some basic tests

import unittest
from apistar import http
from utils.cookie_sessions import init_cookie_session


class CookieSessions(unittest.TestCase):
    def setUp(self):
        self.response_headers = http.ResponseHeaders()

    def test_new_session_is_empty(self):
        with init_cookie_session(http.Header(""), self.response_headers) as session:
            assert session.data == {}

    def test_set_session(self):
        with init_cookie_session(http.Header(""), self.response_headers) as session:
            session["user"] = 1
            session["is_active"] = True
        cookie = http.Header(self.response_headers["Set-Cookie"])
        with init_cookie_session(cookie, self.response_headers) as session:
            assert session["user"] == 1
            assert session.get("is_active") is True
        cookie = http.Header(self.response_headers["Set-Cookie"])
        with init_cookie_session(cookie, self.response_headers) as session:
            assert session.data == {"user": 1, "is_active": True}

    def test_unset_session(self):
        with init_cookie_session(http.Header(""), self.response_headers) as session:
            session["user"] = 1
            session["is_active"] = True
        cookie = http.Header(self.response_headers["Set-Cookie"])
        with init_cookie_session(cookie, self.response_headers) as session:
            del session["user"]
            assert session.get("is_active") is True
        cookie = http.Header(self.response_headers["Set-Cookie"])
        with init_cookie_session(cookie, self.response_headers) as session:
            assert "user" not in session
            assert session.get("is_active") is True
            del session["is_active"]
        cookie = http.Header(self.response_headers["Set-Cookie"])
        with init_cookie_session(cookie, self.response_headers) as session:
            assert session.data == {}

    def test_session_modification(self):
        with init_cookie_session(http.Header(""), self.response_headers) as session:
            session["user"] = 1
        modified_session = self.response_headers["Set-Cookie"].replace("a", "e")
        cookie = http.Header(modified_session)
        with init_cookie_session(cookie, self.response_headers) as session:
            assert session.data == {}