generated from FOSS/Bot-Template
87 lines
3.2 KiB
Python
87 lines
3.2 KiB
Python
# - *- coding: utf- 8 - *-
|
|
import time
|
|
from typing import Any, Awaitable, Callable, Dict, Optional, Union
|
|
|
|
from aiogram import BaseMiddleware
|
|
from aiogram.dispatcher.flags import get_flag
|
|
from aiogram.types import CallbackQuery, Message, TelegramObject, User
|
|
from cachetools import TTLCache
|
|
|
|
from tgbot.data.config import BOT_THROTTLE_RATE
|
|
|
|
|
|
# Простая защита от спама
|
|
class ThrottlingMiddleware(BaseMiddleware):
|
|
def __init__(self, default_rate: Union[int, float] = BOT_THROTTLE_RATE) -> None:
|
|
# Базовая задержка между сообщениями
|
|
self.default_rate = default_rate
|
|
|
|
self.message_users = TTLCache(maxsize=10_000, ttl=600)
|
|
self.callback_users = TTLCache(maxsize=10_000, ttl=600)
|
|
|
|
async def __call__(
|
|
self,
|
|
handler: Callable[[TelegramObject, Dict[str, Any]], Awaitable[Any]],
|
|
event: TelegramObject,
|
|
data,
|
|
):
|
|
# Если юзер спамит, постепенно увеличиваем паузу
|
|
this_user: Optional[User] = data.get("event_from_user")
|
|
|
|
if this_user is None:
|
|
return await handler(event, data)
|
|
|
|
flag_rate = get_flag(data, "rate")
|
|
rate = float(self.default_rate if flag_rate is None else flag_rate)
|
|
|
|
if rate == 0:
|
|
return await handler(event, data)
|
|
|
|
now_time = time.monotonic()
|
|
bucket = self._get_bucket(event)
|
|
user_key = this_user.id
|
|
|
|
if user_key not in bucket:
|
|
bucket[user_key] = {
|
|
"last_throttled": now_time,
|
|
"count_throttled": 0,
|
|
"now_rate": rate,
|
|
}
|
|
|
|
return await handler(event, data)
|
|
else:
|
|
if (
|
|
now_time - bucket[user_key]["last_throttled"]
|
|
>= bucket[user_key]["now_rate"]
|
|
):
|
|
bucket.pop(user_key)
|
|
|
|
return await handler(event, data)
|
|
else:
|
|
bucket[user_key]["last_throttled"] = now_time
|
|
bucket[user_key]["count_throttled"] += 1
|
|
|
|
if bucket[user_key]["count_throttled"] == 1:
|
|
bucket[user_key]["now_rate"] = rate + 2
|
|
await self._warn_user(event)
|
|
elif bucket[user_key]["count_throttled"] == 2:
|
|
bucket[user_key]["now_rate"] = rate + 3
|
|
else:
|
|
bucket[user_key]["now_rate"] = rate + 5
|
|
|
|
return None
|
|
|
|
# Для сообщений и колбэков держим разные лимиты, чтобы они не мешали друг другу
|
|
def _get_bucket(self, event: TelegramObject) -> TTLCache:
|
|
if isinstance(event, CallbackQuery):
|
|
return self.callback_users
|
|
|
|
return self.message_users
|
|
|
|
# Предупреждаем там, где это возможно для конкретного типа апдейта
|
|
async def _warn_user(self, event: TelegramObject) -> None:
|
|
if isinstance(event, Message):
|
|
await event.reply("<b>❗ Пожалуйста, не спамьте</b>")
|
|
elif isinstance(event, CallbackQuery) or hasattr(event, "answer"):
|
|
await event.answer("❗ Пожалуйста, не спамьте", cache_time=1)
|