drhider/site/session.py

"""
Хранилище сессий в памяти с автоочисткой по TTL.

Использование:
    sid = create_session()
    add_file(sid, filename, content_bytes)
    files = get_files(sid)
    store_result(sid, zip_bytes)
    zip_data = get_result(sid)
    cleanup(sid)
"""

import uuid
import threading
import time
from typing import Dict, List, Tuple, Optional

# ═══════════════════════════════════════════════════════════════════════════
# Данные сессии
# ═══════════════════════════════════════════════════════════════════════════

TTL_SECONDS = 30 * 60  # 30 минут

# Максимальный объём одного файла и суммарный объём файлов в сессии (защита памяти)
MAX_FILE_BYTES = 50 * 1024 * 1024  # 50 MB на один файл
MAX_SESSION_BYTES = 500 * 1024 * 1024  # 500 MB суммарно на сессию

_sessions: Dict[str, dict] = {}
_lock = threading.Lock()


def _start_timer(sid: str):
    """Запустить таймер автоочистки сессии через TTL."""

    def _clean():
        with _lock:
            _sessions.pop(sid, None)

    timer = threading.Timer(TTL_SECONDS, _clean)
    timer.daemon = True
    timer.start()
    return timer


def touch(sid: str):
    """Продлить жизнь сессии: перезапустить TTL-таймер (если сессия существует)."""
    with _lock:
        s = _sessions.get(sid)
        if not s:
            return
        if s.get("timer"):
            s["timer"].cancel()
        s["timer"] = _start_timer(sid)


def pause_ttl(sid: str):
    """Приостановить TTL сессии (во время обработки): сессия живёт, пока идёт воркер."""
    with _lock:
        s = _sessions.get(sid)
        if s and s.get("timer"):
            s["timer"].cancel()
            s["timer"] = None


def resume_ttl(sid: str):
    """Возобновить TTL сессии (после завершения обработки): результат доступен ещё TTL."""
    with _lock:
        s = _sessions.get(sid)
        if not s:
            return
        if s.get("timer"):
            s["timer"].cancel()
        s["timer"] = _start_timer(sid)


def request_cancel(sid: str) -> bool:
    """Запросить мягкое прерывание обработки сессии.

    Returns:
        True если сессия существует и отмена запрошена, False если нет.
    """
    with _lock:
        s = _sessions.get(sid)
        if not s:
            return False
        s["cancel"].set()
        return True


def get_cancel_event(sid: str) -> Optional[threading.Event]:
    """Получить событие отмены сессии (или None, если сессии нет)."""
    with _lock:
        s = _sessions.get(sid)
        return s["cancel"] if s else None


# ═══════════════════════════════════════════════════════════════════════════
# API
# ═══════════════════════════════════════════════════════════════════════════

def create_session() -> str:
    """Создать новую сессию.

    Returns:
        Уникальный идентификатор сессии (UUID).
    """
    sid = uuid.uuid4().hex
    with _lock:
        _sessions[sid] = {
            "files": [],
            "result": None,
            "cancel": threading.Event(),
            "timer": _start_timer(sid),
        }
    return sid


def add_file(sid: str, filename: str, content: bytes) -> bool:
    """Добавить файл в сессию.

    Args:
        sid: Идентификатор сессии
        filename: Имя файла
        content: Бинарное содержимое

    Returns:
        True если сессия существует, False если нет.
    """
    with _lock:
        s = _sessions.get(sid)
        if not s:
            return False
        total = sum(len(c) for _, c in s["files"])
        if total + len(content) > MAX_SESSION_BYTES:
            return False  # превышен суммарный лимит сессии
        s["files"].append((filename, content))
        return True


def get_files(sid: str) -> Optional[List[Tuple[str, bytes]]]:
    """Получить все файлы сессии.

    Args:
        sid: Идентификатор сессии

    Returns:
        [(filename, content), ...] или None если сессия не найдена.
    """
    with _lock:
        s = _sessions.get(sid)
        return list(s["files"]) if s else None


def file_count(sid: str) -> int:
    """Количество файлов в сессии."""
    with _lock:
        s = _sessions.get(sid)
        return len(s["files"]) if s else 0


def store_result(sid: str, zip_data: bytes) -> bool:
    """Сохранить результат обработки.

    Args:
        sid: Идентификатор сессии
        zip_data: ZIP-архив с результатом

    Returns:
        True если сессия существует, False если нет.
    """
    with _lock:
        s = _sessions.get(sid)
        if not s:
            return False
        s["result"] = zip_data
        return True


def get_result(sid: str) -> Optional[bytes]:
    """Получить результат обработки.

    Args:
        sid: Идентификатор сессии

    Returns:
        ZIP-архив или None если сессия не найдена/результат не готов.
    """
    with _lock:
        s = _sessions.get(sid)
        return s["result"] if s else None


def cleanup(sid: str):
    """Удалить сессию."""
    with _lock:
        s = _sessions.pop(sid, None)
        if s and s.get("timer"):
            s["timer"].cancel()


def store_csv(sid: str, csv_str: str) -> bool:
    """Сохранить CSV с таблицей замен.

    Args:
        sid: Идентификатор сессии
        csv_str: Строка CSV

    Returns:
        True если сессия существует, False если нет.
    """
    with _lock:
        s = _sessions.get(sid)
        if not s:
            return False
        s["csv"] = csv_str
        return True


def get_csv(sid: str) -> Optional[str]:
    """Получить CSV с таблицей замен.

    Args:
        sid: Идентификатор сессии

    Returns:
        Строка CSV или None если нет.
    """
    with _lock:
        s = _sessions.get(sid)
        return s.get("csv") if s else None