max/max/sessions.py

import time

from flask import request
from flask.sessions import SessionInterface as FlaskSessionInterface, SessionMixin
from itsdangerous import BadSignature, Signer, want_bytes
from werkzeug.datastructures import CallbackDict

from max.crypto import get_random_string
from max.exceptions import SessionWithoutUser

# in seconds
SESSION_IDLE_TIMEOUT = 30 * 60
SESSION_ABSOLUTE_TIMEOUT = 12 * 60 * 60

# Keys for the session metadata in the session storage
CREATED_AT_KEY = "_created_at"
LAST_SEEN_KEY = "_last_seen"


SESSION_MAX_PER_USER = 10


class SessionStorage:
    def __init__(self):
        self.data = {}

    def __getitem__(self, id):
        return self.data.get(id)

    def __delitem__(self, id):
        self.data.pop(id, None)

    def __setitem__(self, id, data):
        if data.get("user_id") is None:
            raise SessionWithoutUser("Refusing to store a session without user_id")

        if id not in self.data:
            self._limit_user_sessions(data["user_id"])
        self.data[id] = data

    def _limit_user_sessions(self, user_id):
        """Makes room for one new session by removing the least recently used ones"""
        user_ids = [id for id, data in self.data.items() if data["user_id"] == user_id]
        user_ids.sort(key=lambda id: self.data[id][LAST_SEEN_KEY])
        for id in user_ids[: max(0, len(user_ids) - SESSION_MAX_PER_USER + 1)]:
            del self.data[id]

    def remove_user_sessions(self, user_id, except_id=None):
        ids_to_remove = []
        for id, data in self.data.items():
            if data["user_id"] == user_id and id != except_id:
                ids_to_remove.append(id)

        for id in ids_to_remove:
            del self.data[id]


session_storage = SessionStorage()


class Session(CallbackDict, SessionMixin):
    def __init__(self, id=None, data=None, created_at=None, last_seen=None):
        def on_update(self):
            self.modified = True
            if not self.id:
                self.set_new_id()

        self.id = id
        self.created_at = created_at
        self.last_seen = last_seen
        self.modified = False
        CallbackDict.__init__(self, data, on_update)

    def is_empty(self):
        return len(self) == 0

    def is_expired(self, now):
        return (
            now - self.created_at > SESSION_ABSOLUTE_TIMEOUT
            or now - self.last_seen > SESSION_IDLE_TIMEOUT
        )

    def clear(self):
        """Removes all data, including the entry in the session storage"""
        del session_storage[self.id]
        self.id = None
        self.created_at = None
        self.last_seen = None
        super().clear()

    def set_new_id(self):
        self.id = get_random_string(32)

    def regenerate(self):
        """Keeps the data under a new id. The old id becomes invalid."""
        del session_storage[self.id]
        self.set_new_id()
        self.modified = True


class SessionInterface(FlaskSessionInterface):
    def _get_signer(self, app):
        if not app.secret_key:
            return None
        return Signer(app.secret_key, salt="flask-session", key_derivation="hmac")

    def get_cookie_name(self, app):
        return app.config["SESSION_COOKIE_NAME"]

    def make_session_object(self, session_id=None, data=None):
        created_at = last_seen = None
        if data is not None:
            data = dict(data)  # do not modify the entry in the session storage
            created_at = data.pop(CREATED_AT_KEY)
            last_seen = data.pop(LAST_SEEN_KEY)
        return Session(
            session_id, data=data, created_at=created_at, last_seen=last_seen
        )

    def open_session(self, app, request):
        session_id = request.cookies.get(self.get_cookie_name(app))
        if not session_id:
            return self.make_session_object()

        signer = self._get_signer(app)
        if signer is None:
            return None
        try:
            sid_as_bytes = signer.unsign(session_id)
            session_id = sid_as_bytes.decode()
        except BadSignature:
            return self.make_session_object()

        if not isinstance(session_id, str):
            session_id = session_id.decode("utf-8", "strict")

        data = session_storage[session_id]
        if data is None:
            # Never adopt an id chosen by the client (session fixation)
            return self.make_session_object()

        session = self.make_session_object(session_id, data=data)
        if session.is_expired(time.monotonic()):
            del session_storage[session_id]
            return self.make_session_object()
        return session

    def get_cookie_attributes(self, app):
        # Attributes for setting and deleting the cookie must match, because
        # browsers ignore a deletion of a __Host- cookie without Secure and Path.
        return dict(
            domain=self.get_cookie_domain(app),
            path=self.get_cookie_path(app),
            secure=self.get_cookie_secure(app),
            httponly=self.get_cookie_httponly(app),
            samesite=self.get_cookie_samesite(app),
        )

    def save_session(self, app, session, response):
        cookie_name = self.get_cookie_name(app)
        cookie_attributes = self.get_cookie_attributes(app)
        if session.is_empty():
            if session.modified:
                del session_storage[session.id]
            # Only send the deletion header if the client sent a cookie
            if cookie_name in request.cookies:
                response.delete_cookie(cookie_name, **cookie_attributes)
        else:
            expires = self.get_expiration_time(app, session)
            session.last_seen = time.monotonic()
            if session.created_at is None:
                session.created_at = session.last_seen
            try:
                session_storage[session.id] = {
                    **session,
                    CREATED_AT_KEY: session.created_at,
                    LAST_SEEN_KEY: session.last_seen,
                }
            except SessionWithoutUser:
                # Flask saves the session again for the error response
                session.clear()
                raise
            session_id = (
                self._get_signer(app).sign(want_bytes(session.id)).decode("utf-8")
            )
            response.set_cookie(
                cookie_name, session_id, expires=expires, **cookie_attributes
            )