chore: add rate limiting for update commands

Author: New-dev0

swibots/bot_app.py

@@ -21,6 +21,7 @@
 from swibots.api.common.models import User
 from swibots.api.callback import AppBar, AppPage
 from swibots.api.common.events import Event
+from swibots.rate_limiter import AsyncRateLimiter
 from swibots.utils import (
     DownloadProgress,
     IOClient,
@@ -81,8 +82,10 @@ def __init__(
         if email and password:
             auth_result = self.auth_service.login(email, password)
             token = auth_result.access_token
+
         if not token:
             raise TokenInvalidError(f"'token' for the bot can't be '{token}'")
+
         self.token = token
         self._user_type = Bot
         self._bot_info: BotInfo | None = None
@@ -95,14 +98,17 @@ def __init__(
         self._bot_description = bot_description
         self.auto_update_bot = auto_update_bot
         self._loop = loop or asyncio.get_event_loop()
-        try:
-            self.user = self.auth_service.get_me_sync(user_type=self._user_type)
-            self.name = self.user.name
-        except Exception as er:
-            logging.error(er)
-            # In case, if issues with auth service
-            self.user = self._user_type(app=self)
-            self.name = ""
+
+        self.user = self.auth_service.get_me_sync(user_type=self._user_type)
+        self.name = self.user.name
+
+        self.rate_limit = AsyncRateLimiter(storage_file=f"{self.user.id}.pkl")
+
+        # Add rate limit functions
+        self.update_bot_commands = self.rate_limit.limit(
+            "update_bot_commands", 10, 15 * 60
+        )(self.update_bot_commands)
+
         if app_bar is None:
             app_bar = AppBar(self.name, left_icon=self.user.imageurl)
         self.app_bar = app_bar
@@ -253,8 +259,9 @@ async def update_bot_commands(self):
             commands=commands,
             preview=self._app_preview,
         )
-
+        log.info("Updating bot commands...")
         self._bot_info = await self.update_bot_info(self._bot_info)
+        log.info("Bot commands updated")
 
     async def _on_chat_service_start(self, _):
         await self.chat_service.subscribe_to_notifications(callback=self.on_chat_event)
@@ -459,5 +466,8 @@ def run(self, task: Callable = None):
                 log.exception(e)
                 run(self.stop())
 
+    def limit(self, key: str, times: int, seconds: int):
+        pass
+
 
 BotApp = Client

swibots/rate_limiter.py

@@ -0,0 +1,44 @@
+import os, pickle, time, aiofiles
+
+
+class AsyncRateLimiter:
+    def __init__(self, storage_file='rate_limit_data.pkl'):
+        self.storage_file = storage_file
+        self.rate_limit_data = {}
+    
+    async def _load_data(self):
+        if os.path.exists(self.storage_file):
+            async with aiofiles.open(self.storage_file, 'rb') as f:
+                content = await f.read()
+                self.rate_limit_data = pickle.loads(content)
+    
+    async def _save_data(self):
+        async with aiofiles.open(self.storage_file, 'wb') as f:
+            await f.write(pickle.dumps(self.rate_limit_data))
+    
+    def limit(self, key: str, times: int, seconds: int):
+        def decorator(func):
+            async def async_wrapper(*args, **kwargs):
+                current_time = time.time()
+
+                await self._load_data()
+                
+                if key not in self.rate_limit_data:
+                    self.rate_limit_data[key] = []
+                
+                self.rate_limit_data[key] = [t for t in self.rate_limit_data[key] if current_time - t < seconds]
+                
+                if len(self.rate_limit_data[key]) >= times:
+                    raise Exception(f"Rate limit exceeded: {times} calls in {seconds} seconds")
+                
+                # Record the call
+                self.rate_limit_data[key].append(current_time)
+                await self._save_data()
+                return await func(*args, **kwargs)
+            
+            return async_wrapper
+        return decorator
+
+
+class RateLimitException(Exception):
+    pass