houseplan-card/scripts/support-relay/hp_relay/ratelimit.py

"""Частотные ограничения без хранения сырых адресов (§9.2 ТЗ).

Адрес источника нигде не сохраняется: он превращается в HMAC от секрета узла и
СЕГОДНЯШНЕЙ даты. Ключ живёт максимум сутки и не позволяет связать обращения
разных дней между собой; секрет узла генерируется при первом старте и лежит
рядом со спулом с правами 0600.
"""

from __future__ import annotations

import hmac
import json
import os
import threading
import time
from dataclasses import dataclass
from hashlib import sha256
from pathlib import Path

from . import config

HOUR = 3600
DAY = 24 * 3600


class RateLimited(Exception):
    """Источник или узел исчерпал лимит."""


def _now() -> float:
    return time.time()


def node_secret(spool: Path) -> bytes:
    path = spool / "node.secret"
    if path.exists():
        return path.read_bytes()
    spool.mkdir(parents=True, exist_ok=True)
    secret = os.urandom(32)
    tmp = path.with_suffix(".tmp")
    tmp.write_bytes(secret)
    tmp.chmod(0o600)
    tmp.replace(path)
    return secret


def source_key(secret: bytes, address: str, now: float | None = None) -> str:
    """Дневной непрозрачный ключ источника: сырой адрес не возвращается никогда."""
    day = time.strftime("%Y-%m-%d", time.gmtime(_now() if now is None else now))
    return hmac.new(secret, f"{day}|{address}".encode("utf-8"), sha256).hexdigest()[:32]


@dataclass
class _Bucket:
    stamps: list[float]

    def prune(self, now: float) -> None:
        self.stamps = [stamp for stamp in self.stamps if now - stamp < DAY]


class Limiter:
    def __init__(self, spool: Path) -> None:
        self._dir = spool / "rate"
        self._dir.mkdir(parents=True, exist_ok=True)
        self._lock = threading.Lock()

    def _path(self, key: str) -> Path:
        return self._dir / f"{key}.json"

    def _load(self, key: str) -> _Bucket:
        path = self._path(key)
        if not path.exists():
            return _Bucket([])
        try:
            return _Bucket(list(json.loads(path.read_text(encoding="utf-8"))))
        except (OSError, ValueError):
            return _Bucket([])

    def _save(self, key: str, bucket: _Bucket) -> None:
        path = self._path(key)
        tmp = path.with_suffix(".tmp")
        tmp.write_text(json.dumps(bucket.stamps), encoding="utf-8")
        tmp.chmod(0o600)
        tmp.replace(path)

    def check_and_count(self, key: str, now: float | None = None) -> None:
        """Считает попытку и бросает RateLimited, если лимит исчерпан.

        Попытка считается ДО доставки: иначе отправитель, добивающийся отказа,
        получал бы бесплатные повторы.
        """
        moment = _now() if now is None else now
        with self._lock:
            for name, limit, window in (
                (key, config.RATE_HOURLY, HOUR),
                (key, config.RATE_DAILY, DAY),
                ("_global", config.RATE_GLOBAL_HOURLY, HOUR),
            ):
                bucket = self._load(name)
                bucket.prune(moment)
                recent = [stamp for stamp in bucket.stamps if moment - stamp < window]
                if len(recent) >= limit:
                    raise RateLimited(name)
            for name in (key, "_global"):
                bucket = self._load(name)
                bucket.prune(moment)
                bucket.stamps.append(moment)
                self._save(name, bucket)

    def purge(self, now: float | None = None) -> int:
        """Удаляет ключи старше суток. Возвращает число удалённых файлов."""
        moment = _now() if now is None else now
        removed = 0
        with self._lock:
            for path in self._dir.glob("*.json"):
                try:
                    stamps = json.loads(path.read_text(encoding="utf-8"))
                except (OSError, ValueError):
                    stamps = []
                if not stamps or moment - max(stamps) >= config.RATE_TTL_SECONDS:
                    path.unlink(missing_ok=True)
                    removed += 1
        return removed