from __future__ import annotations import time from collections import defaultdict, deque from collections.abc import Hashable from fastapi import HTTPException, Request, status from sqlalchemy.ext.asyncio import AsyncSession from app.core.config import settings from app.models.user import User BucketKey = tuple[str, Hashable] _buckets: dict[BucketKey, deque[float]] = defaultdict(deque) _redis_client = None def reset_rate_limit_state() -> None: _buckets.clear() async def check_rate_limit( *, scope: str, limit: int, window_seconds: int, request: Request | None = None, user: User | None = None, session: AsyncSession | None = None, ) -> None: identifiers: list[Hashable] = [] if user is not None: identifiers.append(f"user:{user.id}") identifiers.append(f"telegram:{user.telegram_id}") if request is not None and request.client is not None: identifiers.append(f"ip:{request.client.host}") if not identifiers: identifiers.append("anonymous") if settings.redis_url: allowed = await check_redis_rate_limit(scope, identifiers, limit, window_seconds) if not allowed: await log_rate_limit_event( session, scope=scope, identifier="redis", user=user, request=request, ) raise_rate_limit(scope, window_seconds) return now = time.monotonic() for identifier in identifiers: key = (scope, identifier) bucket = _buckets[key] while bucket and now - bucket[0] > window_seconds: bucket.popleft() if len(bucket) >= limit: await log_rate_limit_event( session, scope=scope, identifier=str(identifier), user=user, request=request, ) raise_rate_limit(scope, window_seconds) for identifier in identifiers: _buckets[(scope, identifier)].append(now) def raise_rate_limit(scope: str, window_seconds: int) -> None: raise HTTPException( status_code=status.HTTP_429_TOO_MANY_REQUESTS, detail={ "code": "rate_limit_exceeded", "message": "Слишком много запросов. Попробуйте чуть позже.", "scope": scope, "retry_after_seconds": window_seconds, }, ) async def get_redis_client(): global _redis_client if _redis_client is not None: return _redis_client try: from redis.asyncio import Redis except ImportError: return None _redis_client = Redis.from_url(settings.redis_url, encoding="utf-8", decode_responses=True) return _redis_client async def check_redis_rate_limit( scope: str, identifiers: list[Hashable], limit: int, window_seconds: int, ) -> bool: client = await get_redis_client() if client is None: return True now_window = int(time.time() // window_seconds) keys = [f"rl:{scope}:{identifier}:{now_window}" for identifier in identifiers] pipe = client.pipeline() for key in keys: pipe.incr(key) pipe.expire(key, window_seconds * 2) results = await pipe.execute() counts = [int(results[index]) for index in range(0, len(results), 2)] return all(count <= limit for count in counts) async def log_rate_limit_event( session: AsyncSession | None, *, scope: str, identifier: str, user: User | None = None, request: Request | None = None, ) -> None: client_host = request.client.host if request and request.client else None user_agent = request.headers.get("user-agent") if request else None metadata = { "scope": scope, "identifier": identifier, "telegram_id": user.telegram_id if user else None, "user_id": user.id if user else None, "ip": client_host, } if session is None: from app.db.session import async_session_factory async with async_session_factory() as event_session: await persist_rate_limit_event( event_session, scope=scope, identifier=identifier, user=user, client_host=client_host, user_agent=user_agent, metadata=metadata, ) return await persist_rate_limit_event( session, scope=scope, identifier=identifier, user=user, client_host=client_host, user_agent=user_agent, metadata=metadata, ) async def persist_rate_limit_event( event_session: AsyncSession, *, scope: str, identifier: str, user: User | None, client_host: str | None, user_agent: str | None, metadata: dict, ) -> None: from app.models.car import AuditLog from app.services.admin_notifications import create_admin_notification try: event_session.add( AuditLog( actor_user_id=user.id if user else None, actor_role=user.platform_role if user else "system", action="rate_limit.exceeded", target_type=scope, target_id=identifier[:80], metadata_json=metadata, ip=client_host, user_agent=user_agent[:256] if user_agent else None, ) ) await create_admin_notification( event_session, event_type="rate_limit_exceeded", title="Rate limit exceeded", body=f"Scope: {scope}\nIdentifier: {identifier}", entity_type="user" if user else "system", entity_id=user.id if user else scope, severity="warning", idempotency_key=f"rate_limit:{scope}:{identifier}:{int(time.time() // max(60, 1))}", metadata=metadata, ) await event_session.commit() except Exception: await event_session.rollback()