""" Rate limiter для ограничения частоты запросов пользователей. Обеспечивает защиту от спама и злоупотреблений. Особенности: - Sliding window алгоритм для точного учета запросов во времени - Поддержка ролевой системы (администраторы не ограничены) - Дифференцированные лимиты для пользователей разных рангов: * Новые пользователи (Рядовой, Ефрейтор): 5 запросов/минуту * Обычные пользователи: 10 запросов/минуту * Администраторы: без ограничений - Потокобезопасная реализация с автоматической очисткой старых записей - Интеграция с системой рангов для определения статуса пользователя """ import time import threading from typing import Dict, Optional, Tuple from collections import defaultdict from .permissions import UserRole class RateLimitExceeded(Exception): """Исключение при превышении лимита запросов""" def __init__(self, retry_after: int): self.retry_after = retry_after super().__init__(f"Rate limit exceeded. Retry after {retry_after} seconds") class RateLimiter: """ Rate limiter на основе sliding window алгоритма. Отслеживает запросы пользователей в скользящем окне времени. Основные возможности: - Проверка лимитов без записи запроса (check_limit) - Запись запроса с проверкой лимита (record_request) - Получение количества оставшихся запросов (get_remaining_requests) - Получение статистики использования (get_stats) - Автоматическая очистка устаревших записей - Поддержка ранговой системы для дифференциации лимитов Лимиты по умолчанию: - Новые пользователи (Рядовой/Ефрейтор): 5 запросов в минуту - Обычные пользователи: 10 запросов в минуту - Администраторы: без ограничений """ def __init__(self, window_size: int = 60, max_requests: int = 10, cleanup_interval: int = 300): """ Инициализация rate limiter. Args: window_size: Размер окна в секундах (по умолчанию 60 секунд = 1 минута) max_requests: Максимальное количество запросов в окне (по умолчанию 10) cleanup_interval: Интервал очистки старых записей в секундах (по умолчанию 5 минут) """ self.window_size = window_size self.max_requests = max_requests self.cleanup_interval = cleanup_interval # Структура: user_id -> список временных меток запросов self.requests: Dict[int, list] = defaultdict(list) # Блокировка для потокобезопасности self.lock = threading.RLock() # Настройки для новых пользователей (более строгие ограничения) self.new_user_max_requests = 5 # 5 запросов в минуту для новых пользователей # Исключения для ролей (администраторы не ограничены) self.exempt_roles = {UserRole.ADMIN, UserRole.SUPER_ADMIN} # Запуск автоматической очистки self._start_cleanup_thread() def _start_cleanup_thread(self): """Запуск фонового потока для очистки старых записей""" def cleanup_worker(): while True: time.sleep(self.cleanup_interval) self._cleanup_old_requests() thread = threading.Thread(target=cleanup_worker, daemon=True) thread.start() def _cleanup_old_requests(self): """Очистка старых временных меток запросов""" current_time = time.time() cutoff_time = current_time - self.window_size with self.lock: for user_id in list(self.requests.keys()): # Оставляем только свежие запросы self.requests[user_id] = [ timestamp for timestamp in self.requests[user_id] if timestamp > cutoff_time ] # Удаляем пустые записи if not self.requests[user_id]: del self.requests[user_id] def _get_max_requests_for_user(self, user_role: Optional[UserRole], is_new_user: bool = False) -> int: """ Определение максимального количества запросов для пользователя. Args: user_role: Роль пользователя is_new_user: Является ли пользователь новым Returns: Максимальное количество запросов в окне """ # Администраторы не ограничены if user_role in self.exempt_roles: return float('inf') # Новые пользователи имеют более строгие ограничения if is_new_user: return self.new_user_max_requests return self.max_requests def _is_new_user(self, user_id: int, registration_time: Optional[float] = None, user_rank: Optional[str] = None) -> bool: """ Определение, является ли пользователь новым на основе ранга или количества запросов. Логика определения новых пользователей: 1. Если указано время регистрации - проверяем, прошло ли менее 24 часов 2. Если известен ранг - проверяем, является ли он начальным ("Рядовой", "Ефрейтор") 3. Иначе проверяем количество предыдущих запросов (< max_requests) Args: user_id: ID пользователя registration_time: Время регистрации (если известно) user_rank: Ранг пользователя (если известен) Returns: True если пользователь новый (низкий ранг или мало запросов) """ # Если время регистрации известно, проверяем его if registration_time: return (time.time() - registration_time) < 86400 # 24 часа # Если ранг известен, считаем нового пользователями с низкими рангами if user_rank: # Новички: Рядовой и Ефрейтор считаются новыми пользователями low_ranks = ["Рядовой", "Ефрейтор"] return user_rank in low_ranks # Иначе проверяем по количеству предыдущих запросов # Если запросов меньше максимального лимита, считаем пользователя новым with self.lock: return len(self.requests.get(user_id, [])) < self.max_requests async def check_limit(self, user_id: int, user_role: Optional[UserRole] = None, registration_time: Optional[float] = None, user_rank: Optional[str] = None) -> Tuple[bool, Optional[int]]: """ Проверка, не превышен ли лимит запросов. Только проверяет, не добавляет запрос в счетчик. Args: user_id: ID пользователя user_role: Роль пользователя registration_time: Время регистрации пользователя user_rank: Ранг пользователя Returns: Кортеж (разрешено: bool, время_ожидания: Optional[int]) """ current_time = time.time() with self.lock: # Определяем лимиты для пользователя is_new_user = self._is_new_user(user_id, registration_time, user_rank) max_requests = self._get_max_requests_for_user(user_role, is_new_user) # Администраторы всегда разрешены if max_requests == float('inf'): return True, None # Очищаем старые запросы для этого пользователя cutoff_time = current_time - self.window_size user_requests = self.requests[user_id] user_requests[:] = [t for t in user_requests if t > cutoff_time] # Проверяем лимит if len(user_requests) >= max_requests: # Вычисляем время до следующего разрешенного запроса oldest_request = min(user_requests) retry_after = int(self.window_size - (current_time - oldest_request)) return False, max(1, retry_after) # Возвращаем разрешение без добавления запроса return True, None async def record_request(self, user_id: int, user_role: Optional[UserRole] = None, registration_time: Optional[float] = None, user_rank: Optional[str] = None) -> None: """ Запись факта выполнения запроса (для случаев, когда проверка уже пройдена). Args: user_id: ID пользователя user_role: Роль пользователя registration_time: Время регистрации пользователя user_rank: Ранг пользователя """ current_time = time.time() with self.lock: # Определяем лимиты для пользователя is_new_user = self._is_new_user(user_id, registration_time, user_rank) max_requests = self._get_max_requests_for_user(user_role, is_new_user) # Администраторы всегда разрешены if max_requests == float('inf'): return # Очищаем старые запросы для этого пользователя cutoff_time = current_time - self.window_size user_requests = self.requests[user_id] user_requests[:] = [t for t in user_requests if t > cutoff_time] # Проверяем лимит if len(user_requests) >= max_requests: # Вычисляем время до следующего разрешенного запроса oldest_request = min(user_requests) retry_after = int(self.window_size - (current_time - oldest_request)) raise RateLimitExceeded(max(1, retry_after)) # Добавляем текущий запрос user_requests.append(current_time) def get_remaining_requests(self, user_id: int, user_role: Optional[UserRole] = None, registration_time: Optional[float] = None, user_rank: Optional[str] = None) -> int: """ Получение количества оставшихся запросов в текущем окне. Args: user_id: ID пользователя user_role: Роль пользователя registration_time: Время регистрации пользователя user_rank: Ранг пользователя Returns: Количество оставшихся запросов """ with self.lock: is_new_user = self._is_new_user(user_id, registration_time, user_rank) max_requests = self._get_max_requests_for_user(user_role, is_new_user) if max_requests == float('inf'): return 999 # Для администраторов возвращаем большое число # Очищаем старые запросы перед подсчетом current_time = time.time() cutoff_time = current_time - self.window_size user_requests = self.requests[user_id] user_requests[:] = [t for t in user_requests if t > cutoff_time] current_requests = len(user_requests) return max(0, max_requests - current_requests) def get_stats(self) -> Dict[str, any]: """ Получение статистики rate limiter. Returns: Словарь со статистикой """ with self.lock: total_users = len(self.requests) total_requests = sum(len(requests) for requests in self.requests.values()) return { 'total_users_tracked': total_users, 'total_requests_in_window': total_requests, 'window_size_seconds': self.window_size, 'default_max_requests': self.max_requests, 'new_user_max_requests': self.new_user_max_requests } # Глобальный экземпляр rate limiter # Настройки согласно требованиям: # - 10 команд/минуту для обычных пользователей # - 5 команд/минуту для новых пользователей (Рядовой, Ефрейтор) # - Без ограничений для администраторов rate_limiter = RateLimiter(window_size=60, max_requests=10)