diff --git a/hikka/modules/api_protection.py b/hikka/modules/api_protection.py index 475e56f..7702b6a 100644 --- a/hikka/modules/api_protection.py +++ b/hikka/modules/api_protection.py @@ -45,19 +45,12 @@ GROUPS = [ CONSTRUCTORS = { - (lambda x: x[0].lower() + x[1:])( - method.__class__.__name__.rsplit("Request", 1)[0] - ): method.CONSTRUCTOR_ID - for method in utils.array_sum( - [ - [ - method - for method in dir(getattr(functions, group)) - if isinstance(method, TLRequest) - ] - for group in GROUPS - ] - ) + (entity_name[0].lower() + entity_name[1:]).rsplit("Request", 1)[0]: getattr(cur_entity, "CONSTRUCTOR_ID") + for group in GROUPS + for entity_name in dir(getattr(functions, group)) + if hasattr((cur_entity := getattr(getattr(functions, group), entity_name)), "__bases__") + and TLRequest in cur_entity.__bases__ + and hasattr(cur_entity, "CONSTRUCTOR_ID") } @@ -101,17 +94,15 @@ class APIRatelimiterMod(loader.Module): "importChatInvite", ] ), - on_change=lambda: self._client.forbid_constructors( - map( - lambda x: CONSTRUCTORS[x], - self.config["forbidden_constructors"], - ) - ), + on_change=self.on_forbidden_methods_update ), ) async def client_ready(self): asyncio.ensure_future(self._install_protection()) + + async def on_forbidden_methods_update(self): + self._client.forbid_constructors(list(map(lambda x: CONSTRUCTORS[x], self.config['forbidden_methods'], ))) async def _install_protection(self): await asyncio.sleep(30) # Restart lock diff --git a/hikka/tl_cache.py b/hikka/tl_cache.py index 7891a45..a77a935 100644 --- a/hikka/tl_cache.py +++ b/hikka/tl_cache.py @@ -685,23 +685,25 @@ class CustomTelegramClient(TelegramClient): flood_sleep_threshold, ) + def _internal_forbid_ctor(self, constructors: list): + self._forbidden_constructors.extend(constructors) + self._forbidden_constructors = list(set(self._forbidden_constructors)) + def forbid_constructor(self, constructor: int): """ Forbids the given constructor to be called :param constructor: Constructor id to forbid """ - self._forbidden_constructors.extend([constructor]) - self._forbidden_constructors = list(set(self._forbidden_constructors)) + self._internal_forbid_ctor([constructor]) def forbid_constructors(self, constructors: list): """ Forbids the given constructors to be called. - All existing forbidden constructors will be removed :param constructors: Constructor ids to forbid """ - self._forbidden_constructors = list(set(constructors)) + self._internal_forbid_ctor(constructors) def _handle_update( self: "CustomTelegramClient",