Refactor to add slowmode delay for messages

Refactor to add slowmode delay for messages

anonflow/bot/middleware/global_slowmode.py

@@ -1,49 +1,52 @@
 import asyncio
+import time
+from typing import Dict, List, Optional
 
 from aiogram import BaseMiddleware
-from aiogram.types import Message
+from aiogram.types import ChatIdUnion, Message
 
 from anonflow.bot.utils.template_renderer import TemplateRenderer
 
+from .utils import extract_message
+
 
 class GlobalSlowmodeMiddleware(BaseMiddleware):
-    def __init__(self,
-                 delay: float,
-                 template_renderer: TemplateRenderer):
+    def __init__(self, delay: float, template_renderer: TemplateRenderer, allowed_chat_ids: Optional[List[ChatIdUnion]] = []):
         super().__init__()
 
         self.delay = delay
         self.renderer = template_renderer
+        self.allowed_chat_ids = allowed_chat_ids
 
+        self.user_times: Dict[int, float] = {}
         self.lock = asyncio.Lock()
 
-    def _extract_message(self, event):
-        if isinstance(event, Message):
-            return event
-        
-        msg = getattr(event, "message", None)
-        if isinstance(msg, Message):
-            return msg
-
-        return None
-
     async def __call__(self, handler, event, data):
-        message = self._extract_message(event)
+        message = extract_message(event)
 
-        if message:
+        if isinstance(message, Message) and message.chat.id not in self.allowed_chat_ids:
             text = message.text or message.caption
             if text and text.startswith("/"):
                 return await handler(event, data)
 
             if self.lock.locked():
+                start_time = self.user_times.get(message.chat.id) or 0
+                current_time = time.monotonic()
+
                 await message.answer(
-                    await self.renderer.render("messages/users/send/busy.j2", message=message)
+                    await self.renderer.render(
+                        "messages/users/send/busy.j2",
+                        message=message,
+                        remaining=round(self.delay - (current_time - start_time)) if start_time else None
+                    )
                 )
                 return
 
             async with self.lock:
                 result = await handler(event, data)
-
+                self.user_times[message.chat.id] = time.monotonic()
                 await asyncio.sleep(self.delay)
 
-                return result
+            return result
+
+        return await handler(event, data)

anonflow/bot/middleware/user_slowmode.py

@@ -0,0 +1,60 @@
+import asyncio
+import time
+from typing import Dict, Optional, List
+
+from aiogram import BaseMiddleware
+from aiogram.types import ChatIdUnion, Message
+
+from anonflow.bot.utils.template_renderer import TemplateRenderer
+
+from .utils import extract_message
+
+
+class UserSlowmodeMiddleware(BaseMiddleware):
+    def __init__(self, delay: float, template_renderer: TemplateRenderer, allowed_chat_ids: Optional[List[ChatIdUnion]] = []):
+        super().__init__()
+
+        self.delay = delay
+        self.renderer = template_renderer
+        self.allowed_chat_ids = allowed_chat_ids
+
+        self.user_times: Dict[int, float] = {}
+        self.user_locks: Dict[int, asyncio.Lock] = {}
+
+        self.lock = asyncio.Lock()
+
+    async def __call__(self, handler, event, data):
+        message = extract_message(event)
+
+        if isinstance(message, Message) and message.chat.id not in self.allowed_chat_ids:
+            text = message.text or message.caption
+            if text and text.startswith("/"):
+                return await handler(event, data)
+
+            async with self.lock:
+                user_lock = self.user_locks.setdefault(message.chat.id, asyncio.Lock())
+
+            if user_lock.locked():
+                start_time = self.user_times.get(message.chat.id) or 0
+                current_time = time.monotonic()
+
+                await message.answer(
+                    await self.renderer.render(
+                        "messages/users/send/busy.j2",
+                        message=message,
+                        remaining=round(self.delay - (current_time - start_time)) if start_time else None
+                    )
+                )
+                return
+
+            async with user_lock:
+                result = await handler(event, data)
+                self.user_times[message.chat.id] = time.monotonic()
+                await asyncio.sleep(self.delay)
+
+            async with self.lock:
+                self.user_locks.pop(message.chat.id)
+
+            return result
+
+        return await handler(event, data)