antiflood

Author: denver-code

README.md

@@ -52,6 +52,8 @@ alt="vtffijuUg5Y" width="240" height="180" border="10" /></a>
 - [ ] Analytics
 - [ ] Rate limit
 - [ ] Inline buttons
+- [ ] Improved structure
+  - [ ] Middlewares
 
 ## Changelog
 ### 0.1.0

bot/__main__.py

@@ -5,6 +5,8 @@
 from aiogram import Bot, types
 from aiogram.dispatcher import Dispatcher
 
+from aiogram.contrib.fsm_storage.memory import MemoryStorage
+
 import asyncio
 from aiogram.utils import executor
 from apscheduler.schedulers.asyncio import AsyncIOScheduler
@@ -19,11 +21,24 @@
 from bot.core.scheduler import scheduler_job
 
 from bot.handlers import setup as handlers_setup
+from bot.middlewares.antiflood import ThrottlingMiddleware, rate_limit
 
 
 scheduler = AsyncIOScheduler()
 
-dp = Dispatcher(bot, loop=asyncio.get_event_loop())
+# storage = RedisStorage2(db=2)
+
+# storage = RedisStorage2(
+#     host="redis:/localhost",
+#     port=6379,
+#     db=7,
+# )
+
+dp = Dispatcher(
+    bot,
+    loop=loop,
+    storage=MemoryStorage(),
+)
 
 
 async def startup(dp: Dispatcher):
@@ -63,6 +78,7 @@ async def startup(dp: Dispatcher):
 
 
 @dp.message_handler(commands=["ping"])
+@rate_limit(5, "ping")
 async def ping(message: types.Message):
     await message.answer("Pong!")
 
@@ -73,6 +89,7 @@ async def time(message: types.Message):
 
 
 def main():
+    dp.middleware.setup(ThrottlingMiddleware())
     scheduler.add_job(
         scheduler_job,
         "interval",

bot/middlewares/antiflood.py

@@ -0,0 +1,140 @@
+import asyncio
+
+from aiogram import Dispatcher, types
+from aiogram.dispatcher import DEFAULT_RATE_LIMIT
+from aiogram.dispatcher.handler import CancelHandler, current_handler
+from aiogram.dispatcher.middlewares import BaseMiddleware
+from aiogram.utils.exceptions import Throttled
+
+
+def rate_limit(limit: int, key=None):
+    """
+
+    Decorator for configuring rate limit and key in different functions.
+
+
+    :param limit:
+
+    :param key:
+
+    :return:
+
+    """
+
+    def decorator(func):
+        setattr(func, "throttling_rate_limit", limit)
+
+        if key:
+            setattr(func, "throttling_key", key)
+
+        return func
+
+    return decorator
+
+
+class ThrottlingMiddleware(BaseMiddleware):
+
+    """
+
+    Simple middleware
+
+    """
+
+    def __init__(self, limit=3, key_prefix="antiflood_"):
+        self.rate_limit = limit
+
+        self.prefix = key_prefix
+
+        super(ThrottlingMiddleware, self).__init__()
+
+    async def on_process_message(self, message: types.Message, data: dict):
+        """
+
+        This handler is called when dispatcher receives a message
+
+
+        :param message:
+
+        """
+
+        # Get current handler
+
+        handler = current_handler.get()
+
+        # Get dispatcher from context
+
+        dispatcher = Dispatcher.get_current()
+
+        # If handler was configured, get rate limit and key from handler
+
+        if handler:
+            limit = getattr(handler, "throttling_rate_limit", self.rate_limit)
+
+            key = getattr(
+                handler, "throttling_key", f"{self.prefix}_{handler.__name__}"
+            )
+
+        else:
+            limit = self.rate_limit
+
+            key = f"{self.prefix}_message"
+
+        # Use Dispatcher.throttle method.
+
+        try:
+            await dispatcher.throttle(key, rate=limit)
+
+        except Throttled as t:
+            # Execute action
+
+            await self.message_throttled(message, t)
+
+            # Cancel current handler
+
+            raise CancelHandler()
+
+    async def message_throttled(self, message: types.Message, throttled: Throttled):
+        """
+
+        Notify user only on first exceed and notify about unlocking only on last exceed
+
+
+        :param message:
+
+        :param throttled:
+
+        """
+
+        handler = current_handler.get()
+
+        dispatcher = Dispatcher.get_current()
+
+        if handler:
+            key = getattr(
+                handler, "throttling_key", f"{self.prefix}_{handler.__name__}"
+            )
+
+        else:
+            key = f"{self.prefix}_message"
+
+        # Calculate how many time is left till the block ends
+
+        delta = throttled.rate - throttled.delta
+
+        # Prevent flooding
+
+        if throttled.exceeded_count <= 2:
+            await message.reply("Too many requests! ")
+
+        # Sleep.
+
+        await asyncio.sleep(delta)
+
+        # Check lock status
+
+        thr = await dispatcher.check_key(key)
+
+        # If current message is not last with current key - do not send message
+
+        if thr.exceeded_count == throttled.exceeded_count:
+            await message.reply("Unlocked.")