Update aiogram 3 template

tgbot/middlewares/middleware_throttling.py

@@ -1,55 +1,78 @@
 # - *- coding: utf- 8 - *-
 import time
-from typing import Any, Awaitable, Callable, Dict, Union
+from typing import Any, Awaitable, Callable, Dict, Optional, Union
 
 from aiogram import BaseMiddleware
 from aiogram.dispatcher.flags import get_flag
-from aiogram.types import Message, User
+from aiogram.types import CallbackQuery, Message, TelegramObject, User
 from cachetools import TTLCache
 
+from tgbot.data.config import BOT_THROTTLE_RATE
 
-# Antiflood
+
+# Простая защита от спама
 class ThrottlingMiddleware(BaseMiddleware):
-    def __init__(self, default_rate: Union[int, float] = 0.5) -> None:
+    def __init__(self, default_rate: Union[int, float] = BOT_THROTTLE_RATE) -> None:
+        # Базовая задержка между сообщениями
         self.default_rate = default_rate
 
-        self.users = TTLCache(maxsize=10_000, ttl=600)
+        self.message_users = TTLCache(maxsize=10_000, ttl=600)
+        self.callback_users = TTLCache(maxsize=10_000, ttl=600)
 
-    async def __call__(self, handler: Callable[[Message, Dict[str, Any]], Awaitable[Any]], event: Message, data):
-        this_user: User = data.get("event_from_user")
+    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 get_flag(data, "rate") is not None:
-            self.default_rate = get_flag(data, "rate")
-
-        if self.default_rate == 0:
+        if this_user is None:
             return await handler(event, data)
 
-        if this_user.id not in self.users:
-            self.users[this_user.id] = {
-                'last_throttled': int(time.time()),
+        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': self.default_rate,
+                'now_rate': rate,
             }
 
             return await handler(event, data)
         else:
-            if int(time.time()) - self.users[this_user.id]['last_throttled'] >= self.users[this_user.id]['now_rate']:
-                self.users.pop(this_user.id)
+            if now_time - bucket[user_key]['last_throttled'] >= bucket[user_key]['now_rate']:
+                bucket.pop(user_key)
 
                 return await handler(event, data)
             else:
-                self.users[this_user.id]['last_throttled'] = int(time.time())
+                bucket[user_key]['last_throttled'] = now_time
+                bucket[user_key]['count_throttled'] += 1
 
-                if self.users[this_user.id]['count_throttled'] == 0:
-                    self.users[this_user.id]['count_throttled'] = 1
-                    self.users[this_user.id]['now_rate'] = self.default_rate + 2
+                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 await handler(event, data)
-                elif self.users[this_user.id]['count_throttled'] == 1:
-                    self.users[this_user.id]['count_throttled'] = 2
-                    self.users[this_user.id]['now_rate'] = self.default_rate + 3
+                return None
 
-                    await event.reply("<b>❗ Please, do not spam.")
-                elif self.users[this_user.id]['count_throttled'] == 2:
-                    self.users[this_user.id]['count_throttled'] = 3
-                    self.users[this_user.id]['now_rate'] = self.default_rate + 5
+    # Для сообщений и колбэков держим разные лимиты, чтобы они не мешали друг другу
+    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)