106 lines
3.6 KiB
Python
106 lines
3.6 KiB
Python
"""Validation and abuse protection for the realtime command protocol."""
|
||
|
||
from __future__ import annotations
|
||
|
||
from collections import defaultdict, deque
|
||
from dataclasses import dataclass
|
||
from time import monotonic
|
||
from typing import Any
|
||
|
||
|
||
class RealtimeProtocolError(ValueError):
|
||
def __init__(self, code: str, detail: str):
|
||
super().__init__(detail)
|
||
self.code = code
|
||
self.detail = detail
|
||
|
||
|
||
def parse_action(content: Any) -> str:
|
||
if not isinstance(content, dict):
|
||
raise RealtimeProtocolError("invalid_payload", "Ожидался JSON-объект.")
|
||
action = content.get("action")
|
||
if not isinstance(action, str) or not action.strip():
|
||
raise RealtimeProtocolError("action_required", "Не указано действие.")
|
||
action = action.strip()
|
||
if len(action) > 64:
|
||
raise RealtimeProtocolError("invalid_action", "Некорректное действие.")
|
||
return action
|
||
|
||
|
||
def parse_positive_int(value: Any, *, field: str) -> int:
|
||
if isinstance(value, bool):
|
||
raise RealtimeProtocolError("invalid_field", f"Поле {field} некорректно.")
|
||
try:
|
||
parsed = int(value)
|
||
except (TypeError, ValueError) as error:
|
||
raise RealtimeProtocolError(
|
||
"invalid_field",
|
||
f"Поле {field} должно быть положительным целым числом.",
|
||
) from error
|
||
if parsed <= 0:
|
||
raise RealtimeProtocolError(
|
||
"invalid_field",
|
||
f"Поле {field} должно быть положительным целым числом.",
|
||
)
|
||
return parsed
|
||
|
||
|
||
def parse_text(value: Any, *, field: str, max_length: int) -> str:
|
||
if not isinstance(value, str):
|
||
raise RealtimeProtocolError(
|
||
"invalid_field",
|
||
f"Поле {field} должно быть строкой.",
|
||
)
|
||
parsed = value.strip()
|
||
if not parsed:
|
||
raise RealtimeProtocolError(
|
||
"invalid_field",
|
||
f"Поле {field} не может быть пустым.",
|
||
)
|
||
if len(parsed) > max_length:
|
||
raise RealtimeProtocolError(
|
||
"field_too_long",
|
||
f"Поле {field} превышает лимит {max_length} символов.",
|
||
)
|
||
return parsed
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class RateLimit:
|
||
requests: int
|
||
window_seconds: float
|
||
|
||
|
||
class ConnectionRateLimiter:
|
||
"""Small per-connection limiter; infrastructure limits remain a second layer."""
|
||
|
||
DEFAULT = RateLimit(requests=120, window_seconds=60)
|
||
GLOBAL = RateLimit(requests=180, window_seconds=60)
|
||
POLICIES = {
|
||
"typing": RateLimit(requests=30, window_seconds=10),
|
||
"stop_typing": RateLimit(requests=30, window_seconds=10),
|
||
"send_message": RateLimit(requests=30, window_seconds=60),
|
||
"calendar_subscribe": RateLimit(requests=30, window_seconds=60),
|
||
}
|
||
|
||
def __init__(self) -> None:
|
||
self._events: dict[str, deque[float]] = defaultdict(deque)
|
||
self._all_events: deque[float] = deque()
|
||
|
||
@staticmethod
|
||
def _consume(events: deque[float], policy: RateLimit, now: float) -> bool:
|
||
threshold = now - policy.window_seconds
|
||
while events and events[0] <= threshold:
|
||
events.popleft()
|
||
if len(events) >= policy.requests:
|
||
return False
|
||
events.append(now)
|
||
return True
|
||
|
||
def allow(self, action: str) -> bool:
|
||
now = monotonic()
|
||
if not self._consume(self._all_events, self.GLOBAL, now):
|
||
return False
|
||
policy = self.POLICIES.get(action, self.DEFAULT)
|
||
return self._consume(self._events[action], policy, now)
|