src/liquidcode/middleware/session.py - develop-0.4 - liquid-g/liquid-code

src/liquidcode/middleware/session.py
121 строка4 KB
User
ci: настроен Black и flake8 для автоматического форматирования кода
04 июл 2026, 08:14
04 июл 2026, 08:144778ed6
"""
Middleware для управления сессиями через подписанные куки.
"""

import hashlib
import hmac
from typing import Any, Callable, Dict, Optional
from uuid import uuid4

from ..contracts import Request, Response
from ..http import HttpResponse
from .base import Middleware


class SessionStore:
    """Хранилище сессий в памяти (для разработки)."""

    def __init__(self):
        self._data: Dict[str, Dict[str, Any]] = {}

    def get(self, session_id: str) -> Dict[str, Any]:
        return self._data.get(session_id, {})

    def set(self, session_id: str, data: Dict[str, Any]) -> None:
        self._data[session_id] = data

    def delete(self, session_id: str) -> None:
        if session_id in self._data:
            del self._data[session_id]

    def exists(self, session_id: str) -> bool:
        return session_id in self._data


class SessionMiddleware(Middleware):
    """
    Middleware для управления сессиями.
    Создаёт куку с session_id, загружает данные сессии в request,
    и сохраняет изменения после обработки запроса.
    """

    def __init__(
        self,
        cookie_name: str = "session_id",
        secret_key: str = "change_me_in_production",
        max_age: int = 86400,  # 24 часа
        store: Optional[SessionStore] = None,
    ):
        self.cookie_name = cookie_name
        self.secret_key = secret_key.encode("utf-8")
        self.max_age = max_age
        self.store = store or SessionStore()

    def _sign(self, value: str) -> str:
        return hmac.new(
            self.secret_key, value.encode("utf-8"), hashlib.sha256
        ).hexdigest()

    def _verify(self, value: str, signature: str) -> bool:
        expected = self._sign(value)
        return hmac.compare_digest(expected, signature)

    def _decode_cookie(self, cookie_value: str) -> Optional[str]:
        try:
            value, signature = cookie_value.rsplit(".", 1)
            if self._verify(value, signature):
                return value
        except ValueError:
            pass
        return None

    def _encode_cookie(self, session_id: str) -> str:
        return f"{session_id}.{self._sign(session_id)}"

    def __call__(
        self, request: Request, call_next: Callable[[Request], Response]
    ) -> Response:
        # 1. Извлечь session_id из куки
        cookie_header = request.get_headers().get("Cookie", "")
        session_id = None
        for part in cookie_header.split(";"):
            part = part.strip()
            if part.startswith(self.cookie_name + "="):
                raw = part[len(self.cookie_name) + 1 :]
                session_id = self._decode_cookie(raw)
                break

        # 2. Если нет валидной сессии — создаём новую
        if not session_id or not self.store.exists(session_id):
            session_id = uuid4().hex
            self.store.set(session_id, {})

        # 3. Загружаем данные в request (через атрибуты)
        request.set_attr("session", self.store.get(session_id))
        request.set_attr("session_id", session_id)
        request.set_attr("session_modified", False)

        # 4. Вызываем следующий middleware/контроллер
        response = call_next(request)

        # 5. После ответа — обновляем хранилище, если сессия изменилась
        if request.get_attr("session_modified", False):
            session_data = request.get_attr("session")
            self.store.set(session_id, session_data)

        # 6. Проверяем, есть ли уже кука в заголовках ответа
        headers = response.get_headers().copy()
        if "Set-Cookie" not in headers:
            cookie_value = (
                f"{self.cookie_name}={self._encode_cookie(session_id)}; "
                f"Max-Age={self.max_age}; Path=/; HttpOnly; SameSite=Lax"
            )
            headers["Set-Cookie"] = cookie_value
            # Создаём новый HttpResponse с обновлёнными заголовками
            response = HttpResponse(
                content=response.get_content(),
                status=response.get_status(),
                headers=headers,
            )

        return response