loader.py 41 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240
  1. """Registers modules"""
  2. # ©️ Dan Gazizullin, 2021-2023
  3. # This file is a part of Hikka Userbot
  4. # 🌐 https://github.com/hikariatama/Hikka
  5. # You can redistribute it and/or modify it under the terms of the GNU AGPLv3
  6. # 🔑 https://www.gnu.org/licenses/agpl-3.0.html
  7. import asyncio
  8. import contextlib
  9. import copy
  10. import importlib
  11. import importlib.machinery
  12. import importlib.util
  13. import inspect
  14. import logging
  15. import os
  16. import re
  17. import sys
  18. import typing
  19. from functools import wraps
  20. from pathlib import Path
  21. from types import FunctionType, ModuleType
  22. from uuid import uuid4
  23. from hikkatl.tl.tlobject import TLObject
  24. from . import security, utils, validators
  25. from .database import Database
  26. from .inline.core import InlineManager
  27. from .translations import Strings, Translator
  28. from .types import (
  29. Command,
  30. ConfigValue,
  31. CoreOverwriteError,
  32. CoreUnloadError,
  33. DragonModule,
  34. InlineMessage,
  35. JSONSerializable,
  36. Library,
  37. LibraryConfig,
  38. LoadError,
  39. Module,
  40. ModuleConfig,
  41. SelfSuspend,
  42. SelfUnload,
  43. StopLoop,
  44. StringLoader,
  45. get_callback_handlers,
  46. get_commands,
  47. get_inline_handlers,
  48. )
  49. __all__ = [
  50. "Modules",
  51. "InfiniteLoop",
  52. "Command",
  53. "CoreOverwriteError",
  54. "CoreUnloadError",
  55. "DragonModule",
  56. "InlineMessage",
  57. "JSONSerializable",
  58. "Library",
  59. "LibraryConfig",
  60. "LoadError",
  61. "Module",
  62. "SelfSuspend",
  63. "SelfUnload",
  64. "StopLoop",
  65. "StringLoader",
  66. "get_commands",
  67. "get_inline_handlers",
  68. "get_callback_handlers",
  69. "validators",
  70. "Database",
  71. "InlineManager",
  72. "Strings",
  73. "Translator",
  74. "ConfigValue",
  75. "ModuleConfig",
  76. "owner",
  77. "group_owner",
  78. "group_admin_add_admins",
  79. "group_admin_change_info",
  80. "group_admin_ban_users",
  81. "group_admin_delete_messages",
  82. "group_admin_pin_messages",
  83. "group_admin_invite_users",
  84. "group_admin",
  85. "group_member",
  86. "pm",
  87. "unrestricted",
  88. "inline_everyone",
  89. "loop",
  90. ]
  91. logger = logging.getLogger(__name__)
  92. owner = security.owner
  93. # deprecated
  94. sudo = security.sudo
  95. support = security.support
  96. # /deprecated
  97. group_owner = security.group_owner
  98. group_admin_add_admins = security.group_admin_add_admins
  99. group_admin_change_info = security.group_admin_change_info
  100. group_admin_ban_users = security.group_admin_ban_users
  101. group_admin_delete_messages = security.group_admin_delete_messages
  102. group_admin_pin_messages = security.group_admin_pin_messages
  103. group_admin_invite_users = security.group_admin_invite_users
  104. group_admin = security.group_admin
  105. group_member = security.group_member
  106. pm = security.pm
  107. unrestricted = security.unrestricted
  108. inline_everyone = security.inline_everyone
  109. async def stop_placeholder() -> bool:
  110. return True
  111. class Placeholder:
  112. """Placeholder"""
  113. VALID_PIP_PACKAGES = re.compile(
  114. r"^\s*# ?requires:(?: ?)((?:{url} )*(?:{url}))\s*$".format(
  115. url=r"[-[\]_.~:/?#@!$&'()*+,;%<=>a-zA-Z0-9]+"
  116. ),
  117. re.MULTILINE,
  118. )
  119. USER_INSTALL = "PIP_TARGET" not in os.environ and "VIRTUAL_ENV" not in os.environ
  120. class InfiniteLoop:
  121. _task = None
  122. status = False
  123. module_instance = None # Will be passed later
  124. def __init__(
  125. self,
  126. func: FunctionType,
  127. interval: int,
  128. autostart: bool,
  129. wait_before: bool,
  130. stop_clause: typing.Union[str, None],
  131. ):
  132. self.func = func
  133. self.interval = interval
  134. self._wait_before = wait_before
  135. self._stop_clause = stop_clause
  136. self.autostart = autostart
  137. def _stop(self, *args, **kwargs):
  138. self._wait_for_stop.set()
  139. def stop(self, *args, **kwargs):
  140. with contextlib.suppress(AttributeError):
  141. _hikka_client_id_logging_tag = copy.copy( # noqa: F841
  142. self.module_instance.allmodules.client.tg_id
  143. )
  144. if self._task:
  145. logger.debug("Stopped loop for method %s", self.func)
  146. self._wait_for_stop = asyncio.Event()
  147. self.status = False
  148. self._task.add_done_callback(self._stop)
  149. self._task.cancel()
  150. return asyncio.ensure_future(self._wait_for_stop.wait())
  151. logger.debug("Loop is not running")
  152. return asyncio.ensure_future(stop_placeholder())
  153. def start(self, *args, **kwargs):
  154. with contextlib.suppress(AttributeError):
  155. _hikka_client_id_logging_tag = copy.copy( # noqa: F841
  156. self.module_instance.allmodules.client.tg_id
  157. )
  158. if not self._task:
  159. logger.debug("Started loop for method %s", self.func)
  160. self._task = asyncio.ensure_future(self.actual_loop(*args, **kwargs))
  161. else:
  162. logger.debug("Attempted to start already running loop")
  163. async def actual_loop(self, *args, **kwargs):
  164. # Wait for loader to set attribute
  165. while not self.module_instance:
  166. await asyncio.sleep(0.01)
  167. if isinstance(self._stop_clause, str) and self._stop_clause:
  168. self.module_instance.set(self._stop_clause, True)
  169. self.status = True
  170. while self.status:
  171. if self._wait_before:
  172. await asyncio.sleep(self.interval)
  173. if (
  174. isinstance(self._stop_clause, str)
  175. and self._stop_clause
  176. and not self.module_instance.get(self._stop_clause, False)
  177. ):
  178. break
  179. try:
  180. await self.func(self.module_instance, *args, **kwargs)
  181. except StopLoop:
  182. break
  183. except Exception:
  184. logger.exception("Error running loop!")
  185. if not self._wait_before:
  186. await asyncio.sleep(self.interval)
  187. self._wait_for_stop.set()
  188. self.status = False
  189. def __del__(self):
  190. self.stop()
  191. def loop(
  192. interval: int = 5,
  193. autostart: typing.Optional[bool] = False,
  194. wait_before: typing.Optional[bool] = False,
  195. stop_clause: typing.Optional[str] = None,
  196. ) -> FunctionType:
  197. """
  198. Create new infinite loop from class method
  199. :param interval: Loop iterations delay
  200. :param autostart: Start loop once module is loaded
  201. :param wait_before: Insert delay before actual iteration, rather than after
  202. :param stop_clause: Database key, based on which the loop will run.
  203. This key will be set to `True` once loop is started,
  204. and will stop after key resets to `False`
  205. :attr status: Boolean, describing whether the loop is running
  206. """
  207. def wrapped(func):
  208. return InfiniteLoop(func, interval, autostart, wait_before, stop_clause)
  209. return wrapped
  210. MODULES_NAME = "modules"
  211. ru_keys = 'ёйцукенгшщзхъфывапролджэячсмитьбю.Ё"№;%:?ЙЦУКЕНГШЩЗХЪФЫВАПРОЛДЖЭ/ЯЧСМИТЬБЮ,'
  212. en_keys = "`qwertyuiop[]asdfghjkl;'zxcvbnm,./~@#$%^&QWERTYUIOP{}ASDFGHJKL:\"|ZXCVBNM<>?"
  213. BASE_DIR = (
  214. "/data"
  215. if "DOCKER" in os.environ
  216. else os.path.normpath(os.path.join(utils.get_base_dir(), ".."))
  217. )
  218. LOADED_MODULES_DIR = os.path.join(BASE_DIR, "loaded_modules")
  219. LOADED_MODULES_PATH = Path(LOADED_MODULES_DIR)
  220. LOADED_MODULES_PATH.mkdir(parents=True, exist_ok=True)
  221. def translatable_docstring(cls):
  222. """Decorator that makes triple-quote docstrings translatable"""
  223. @wraps(cls.config_complete)
  224. def config_complete(self, *args, **kwargs):
  225. def proccess_decorators(mark: str, obj: str):
  226. nonlocal self
  227. for attr in dir(func_):
  228. if (
  229. attr.endswith("_doc")
  230. and len(attr) == 6
  231. and isinstance(getattr(func_, attr), str)
  232. ):
  233. var = f"strings_{attr.split('_')[0]}"
  234. if not hasattr(self, var):
  235. setattr(self, var, {})
  236. getattr(self, var).setdefault(f"{mark}{obj}", getattr(func_, attr))
  237. for command_, func_ in get_commands(cls).items():
  238. proccess_decorators("_cmd_doc_", command_)
  239. try:
  240. func_.__doc__ = self.strings[f"_cmd_doc_{command_}"]
  241. except AttributeError:
  242. func_.__func__.__doc__ = self.strings[f"_cmd_doc_{command_}"]
  243. for inline_handler_, func_ in get_inline_handlers(cls).items():
  244. proccess_decorators("_ihandle_doc_", inline_handler_)
  245. try:
  246. func_.__doc__ = self.strings[f"_ihandle_doc_{inline_handler_}"]
  247. except AttributeError:
  248. func_.__func__.__doc__ = self.strings[f"_ihandle_doc_{inline_handler_}"]
  249. self.__doc__ = self.strings["_cls_doc"]
  250. return (
  251. self.config_complete._old_(self, *args, **kwargs)
  252. if not kwargs.pop("reload_dynamic_translate", None)
  253. else True
  254. )
  255. config_complete._old_ = cls.config_complete
  256. cls.config_complete = config_complete
  257. for command_, func in get_commands(cls).items():
  258. cls.strings[f"_cmd_doc_{command_}"] = inspect.getdoc(func)
  259. for inline_handler_, func in get_inline_handlers(cls).items():
  260. cls.strings[f"_ihandle_doc_{inline_handler_}"] = inspect.getdoc(func)
  261. cls.strings["_cls_doc"] = inspect.getdoc(cls)
  262. return cls
  263. tds = translatable_docstring # Shorter name for modules to use
  264. def ratelimit(func: Command) -> Command:
  265. """Decorator that causes ratelimiting for this command to be enforced more strictly"""
  266. func.ratelimit = True
  267. return func
  268. def tag(*tags, **kwarg_tags):
  269. """
  270. Tag function (esp. watchers) with some tags
  271. Currently available tags:
  272. • `no_commands` - Ignore all userbot commands in watcher
  273. • `only_commands` - Capture only userbot commands in watcher
  274. • `out` - Capture only outgoing events
  275. • `in` - Capture only incoming events
  276. • `only_messages` - Capture only messages (not join events)
  277. • `editable` - Capture only messages, which can be edited (no forwards etc.)
  278. • `no_media` - Capture only messages without media and files
  279. • `only_media` - Capture only messages with media and files
  280. • `only_photos` - Capture only messages with photos
  281. • `only_videos` - Capture only messages with videos
  282. • `only_audios` - Capture only messages with audios
  283. • `only_docs` - Capture only messages with documents
  284. • `only_stickers` - Capture only messages with stickers
  285. • `only_inline` - Capture only messages with inline queries
  286. • `only_channels` - Capture only messages with channels
  287. • `only_groups` - Capture only messages with groups
  288. • `only_pm` - Capture only messages with private chats
  289. • `no_pm` - Exclude messages with private chats
  290. • `no_channels` - Exclude messages with channels
  291. • `no_groups` - Exclude messages with groups
  292. • `no_inline` - Exclude messages with inline queries
  293. • `no_stickers` - Exclude messages with stickers
  294. • `no_docs` - Exclude messages with documents
  295. • `no_audios` - Exclude messages with audios
  296. • `no_videos` - Exclude messages with videos
  297. • `no_photos` - Exclude messages with photos
  298. • `no_forwards` - Exclude forwarded messages
  299. • `no_reply` - Exclude messages with replies
  300. • `no_mention` - Exclude messages with mentions
  301. • `mention` - Capture only messages with mentions
  302. • `only_reply` - Capture only messages with replies
  303. • `only_forwards` - Capture only forwarded messages
  304. • `startswith` - Capture only messages that start with given text
  305. • `endswith` - Capture only messages that end with given text
  306. • `contains` - Capture only messages that contain given text
  307. • `regex` - Capture only messages that match given regex
  308. • `filter` - Capture only messages that pass given function
  309. • `from_id` - Capture only messages from given user
  310. • `chat_id` - Capture only messages from given chat
  311. • `thumb_url` - Works for inline command handlers. Will be shown in help
  312. • `alias` - Set single alias for a command
  313. • `aliases` - Set multiple aliases for a command
  314. Usage example:
  315. @loader.tag("no_commands", "out")
  316. @loader.tag("no_commands", out=True)
  317. @loader.tag(only_messages=True)
  318. @loader.tag("only_messages", "only_pm", regex=r"^[.] ?hikka$", from_id=659800858)
  319. 💡 These tags can be used directly in `@loader.watcher`:
  320. @loader.watcher("no_commands", out=True)
  321. """
  322. def inner(func: Command) -> Command:
  323. for _tag in tags:
  324. setattr(func, _tag, True)
  325. for _tag, value in kwarg_tags.items():
  326. setattr(func, _tag, value)
  327. return func
  328. return inner
  329. def _mark_method(mark: str, *args, **kwargs) -> typing.Callable[..., Command]:
  330. """
  331. Mark method as a method of a class
  332. """
  333. def decorator(func: Command) -> Command:
  334. setattr(func, mark, True)
  335. for arg in args:
  336. setattr(func, arg, True)
  337. for kwarg, value in kwargs.items():
  338. setattr(func, kwarg, value)
  339. return func
  340. return decorator
  341. def command(*args, **kwargs):
  342. """
  343. Decorator that marks function as userbot command
  344. """
  345. return _mark_method("is_command", *args, **kwargs)
  346. def debug_method(*args, **kwargs):
  347. """
  348. Decorator that marks function as IDM (Internal Debug Method)
  349. :param name: Name of the method
  350. """
  351. return _mark_method("is_debug_method", *args, **kwargs)
  352. def inline_handler(*args, **kwargs):
  353. """
  354. Decorator that marks function as inline handler
  355. """
  356. return _mark_method("is_inline_handler", *args, **kwargs)
  357. def watcher(*args, **kwargs):
  358. """
  359. Decorator that marks function as watcher
  360. """
  361. return _mark_method("is_watcher", *args, **kwargs)
  362. def callback_handler(*args, **kwargs):
  363. """
  364. Decorator that marks function as callback handler
  365. """
  366. return _mark_method("is_callback_handler", *args, **kwargs)
  367. def raw_handler(*updates: TLObject):
  368. """
  369. Decorator that marks function as raw telethon events handler
  370. Use it to prevent zombie-event-handlers, left by unloaded modules
  371. :param updates: Update(-s) to handle
  372. ⚠️ Do not try to simulate behavior of this decorator by yourself!
  373. ⚠️ This feature won't work, if you dynamically declare method with decorator!
  374. """
  375. def inner(func: Command) -> Command:
  376. func.is_raw_handler = True
  377. func.updates = updates
  378. func.id = uuid4().hex
  379. return func
  380. return inner
  381. class Modules:
  382. """Stores all registered modules"""
  383. def __init__(
  384. self,
  385. client: "CustomTelegramClient", # type: ignore # noqa: F821
  386. db: Database,
  387. allclients: list,
  388. translator: Translator,
  389. ):
  390. self._initial_registration = True
  391. self.commands = {}
  392. self.inline_handlers = {}
  393. self.callback_handlers = {}
  394. self.aliases = {}
  395. self.modules = [] # skipcq: PTC-W0052
  396. self.dragon_modules = []
  397. self.libraries = []
  398. self.watchers = []
  399. self._log_handlers = []
  400. self._core_commands = []
  401. self.__approve = []
  402. self.allclients = allclients
  403. self.client = client
  404. self._db = db
  405. self.db = db
  406. self.translator = translator
  407. self.secure_boot = False
  408. asyncio.ensure_future(self._junk_collector())
  409. self.inline = InlineManager(self.client, self._db, self)
  410. self.client.hikka_inline = self.inline
  411. async def _junk_collector(self):
  412. """
  413. Periodically reloads commands, inline handlers, callback handlers and watchers from loaded
  414. modules to prevent zombie handlers
  415. """
  416. while True:
  417. await asyncio.sleep(30)
  418. commands = {}
  419. inline_handlers = {}
  420. callback_handlers = {}
  421. watchers = []
  422. for module in self.modules:
  423. commands.update(module.hikka_commands)
  424. inline_handlers.update(module.hikka_inline_handlers)
  425. callback_handlers.update(module.hikka_callback_handlers)
  426. watchers.extend(module.hikka_watchers.values())
  427. self.commands = commands
  428. self.inline_handlers = inline_handlers
  429. self.callback_handlers = callback_handlers
  430. self.watchers = watchers
  431. logger.debug(
  432. (
  433. "Reloaded %s commands,"
  434. " %s inline handlers,"
  435. " %s callback handlers and"
  436. " %s watchers"
  437. ),
  438. len(self.commands),
  439. len(self.inline_handlers),
  440. len(self.callback_handlers),
  441. len(self.watchers),
  442. )
  443. async def register_all(
  444. self,
  445. mods: typing.Optional[typing.List[str]] = None,
  446. no_external: bool = False,
  447. ) -> typing.List[Module]:
  448. """Load all modules in the module directory"""
  449. external_mods = []
  450. if not mods:
  451. mods = [
  452. os.path.join(utils.get_base_dir(), MODULES_NAME, mod)
  453. for mod in filter(
  454. lambda x: (x.endswith(".py") and not x.startswith("_")),
  455. os.listdir(os.path.join(utils.get_base_dir(), MODULES_NAME)),
  456. )
  457. ]
  458. self.secure_boot = self._db.get(__name__, "secure_boot", False)
  459. external_mods = (
  460. []
  461. if self.secure_boot
  462. else [
  463. (LOADED_MODULES_PATH / mod).resolve()
  464. for mod in filter(
  465. lambda x: (
  466. x.endswith(f"{self.client.tg_id}.py")
  467. and not x.startswith("_")
  468. ),
  469. os.listdir(LOADED_MODULES_DIR),
  470. )
  471. ]
  472. )
  473. loaded = []
  474. loaded += await self._register_modules(mods)
  475. if not no_external:
  476. loaded += await self._register_modules(external_mods, "<file>")
  477. return loaded
  478. async def _register_modules(
  479. self,
  480. modules: list,
  481. origin: str = "<core>",
  482. ) -> typing.List[Module]:
  483. with contextlib.suppress(AttributeError):
  484. _hikka_client_id_logging_tag = copy.copy(self.client.tg_id) # noqa: F841
  485. loaded = []
  486. for mod in modules:
  487. try:
  488. mod_shortname = os.path.basename(mod).rsplit(".py", maxsplit=1)[0]
  489. module_name = f"{__package__}.{MODULES_NAME}.{mod_shortname}"
  490. user_friendly_origin = (
  491. "<core {}>" if origin == "<core>" else "<file {}>"
  492. ).format(module_name)
  493. logger.debug("Loading %s from filesystem", module_name)
  494. spec = importlib.machinery.ModuleSpec(
  495. module_name,
  496. StringLoader(Path(mod).read_text(), user_friendly_origin),
  497. origin=user_friendly_origin,
  498. )
  499. loaded += [await self.register_module(spec, module_name, origin)]
  500. except Exception as e:
  501. logger.exception("Failed to load module %s due to %s:", mod, e)
  502. return loaded
  503. def register_dragon(self, module: ModuleType, instance: DragonModule):
  504. for mod in self.dragon_modules.copy():
  505. if mod.name == instance.name:
  506. logger.debug("Removing dragon module %s for reload", mod.name)
  507. self.unload_dragon(mod)
  508. instance.handlers = []
  509. for name, obj in vars(module).items():
  510. for handler, group in getattr(obj, "handlers", []):
  511. try:
  512. handler = self.client.pyro_proxy.add_handler(handler, group)
  513. instance.handlers.append(handler)
  514. except Exception as e:
  515. logging.exception(
  516. "Can't add handler %s due to %s: %s",
  517. name,
  518. type(e).__name__,
  519. e,
  520. )
  521. self.dragon_modules += [instance]
  522. def unload_dragon(self, instance: DragonModule) -> bool:
  523. for handler in instance.handlers:
  524. try:
  525. self.client.pyro_proxy.remove_handler(*handler)
  526. except Exception as e:
  527. logging.exception(
  528. "Can't remove handler %s due to %s: %s",
  529. handler,
  530. type(e).__name__,
  531. e,
  532. )
  533. if instance in self.dragon_modules:
  534. self.dragon_modules.remove(instance)
  535. return True
  536. return False
  537. async def register_module(
  538. self,
  539. spec: importlib.machinery.ModuleSpec,
  540. module_name: str,
  541. origin: str = "<core>",
  542. save_fs: bool = False,
  543. is_dragon: bool = False,
  544. ) -> typing.Union[Module, typing.Tuple[ModuleType, DragonModule]]:
  545. """Register single module from importlib spec"""
  546. with contextlib.suppress(AttributeError):
  547. _hikka_client_id_logging_tag = copy.copy(self.client.tg_id) # noqa: F841
  548. module = importlib.util.module_from_spec(spec)
  549. sys.modules[module_name] = module
  550. spec.loader.exec_module(module)
  551. if is_dragon:
  552. return module, DragonModule()
  553. ret = None
  554. ret = next(
  555. (
  556. value()
  557. for value in vars(module).values()
  558. if inspect.isclass(value) and issubclass(value, Module)
  559. ),
  560. None,
  561. )
  562. if hasattr(module, "__version__"):
  563. ret.__version__ = module.__version__
  564. if ret is None:
  565. ret = module.register(module_name)
  566. if not isinstance(ret, Module):
  567. raise TypeError(f"Instance is not a Module, it is {type(ret)}")
  568. await self.complete_registration(ret)
  569. ret.__origin__ = origin
  570. cls_name = ret.__class__.__name__
  571. if save_fs:
  572. path = os.path.join(
  573. LOADED_MODULES_DIR,
  574. f"{cls_name}_{self.client.tg_id}.py",
  575. )
  576. if origin == "<string>":
  577. Path(path).write_text(spec.loader.data.decode())
  578. logger.debug("Saved class %s to path %s", cls_name, path)
  579. return ret
  580. def add_aliases(self, aliases: dict):
  581. """Saves aliases and applies them to <core>/<file> modules"""
  582. self.aliases.update(aliases)
  583. for alias, cmd in aliases.items():
  584. self.add_alias(alias, cmd)
  585. def register_raw_handlers(self, instance: Module):
  586. """Register event handlers for a module"""
  587. for name, handler in utils.iter_attrs(instance):
  588. if getattr(handler, "is_raw_handler", False):
  589. self.client.dispatcher.raw_handlers.append(handler)
  590. logger.debug(
  591. "Registered raw handler %s for %s. ID: %s",
  592. name,
  593. instance.__class__.__name__,
  594. handler.id,
  595. )
  596. @property
  597. def _remove_core_protection(self) -> bool:
  598. from . import main
  599. return self._db.get(main.__name__, "remove_core_protection", False)
  600. def register_commands(self, instance: Module):
  601. """Register commands from instance"""
  602. with contextlib.suppress(AttributeError):
  603. _hikka_client_id_logging_tag = copy.copy(self.client.tg_id) # noqa: F841
  604. if instance.__origin__.startswith("<core"):
  605. self._core_commands += list(
  606. map(lambda x: x.lower(), list(instance.hikka_commands))
  607. )
  608. for _command, cmd in instance.hikka_commands.items():
  609. # Restrict overwriting core modules' commands
  610. if (
  611. not self._remove_core_protection
  612. and _command.lower() in self._core_commands
  613. and not instance.__origin__.startswith("<core")
  614. ):
  615. with contextlib.suppress(Exception):
  616. self.modules.remove(instance)
  617. raise CoreOverwriteError(command=_command)
  618. self.commands.update({_command.lower(): cmd})
  619. for alias, cmd in self.aliases.copy().items():
  620. if cmd in instance.hikka_commands:
  621. self.add_alias(alias, cmd)
  622. self.register_inline_stuff(instance)
  623. def register_inline_stuff(self, instance: Module):
  624. for name, func in instance.hikka_inline_handlers.copy().items():
  625. if name.lower() in self.inline_handlers:
  626. if (
  627. hasattr(func, "__self__")
  628. and hasattr(self.inline_handlers[name], "__self__")
  629. and (
  630. func.__self__.__class__.__name__
  631. != self.inline_handlers[name].__self__.__class__.__name__
  632. )
  633. ):
  634. logger.debug(
  635. "Duplicate inline_handler %s of %s",
  636. name,
  637. instance.__class__.__name__,
  638. )
  639. logger.debug(
  640. "Replacing inline_handler %s for %s",
  641. self.inline_handlers[name],
  642. instance.__class__.__name__,
  643. )
  644. self.inline_handlers.update({name.lower(): func})
  645. for name, func in instance.hikka_callback_handlers.copy().items():
  646. if name.lower() in self.callback_handlers and (
  647. hasattr(func, "__self__")
  648. and hasattr(self.callback_handlers[name], "__self__")
  649. and func.__self__.__class__.__name__
  650. != self.callback_handlers[name].__self__.__class__.__name__
  651. ):
  652. logger.debug(
  653. "Duplicate callback_handler %s of %s",
  654. name,
  655. instance.__class__.__name__,
  656. )
  657. self.callback_handlers.update({name.lower(): func})
  658. def unregister_inline_stuff(self, instance: Module, purpose: str):
  659. for name, func in instance.hikka_inline_handlers.copy().items():
  660. if name.lower() in self.inline_handlers and (
  661. hasattr(func, "__self__")
  662. and hasattr(self.inline_handlers[name], "__self__")
  663. and func.__self__.__class__.__name__
  664. == self.inline_handlers[name].__self__.__class__.__name__
  665. ):
  666. del self.inline_handlers[name.lower()]
  667. logger.debug(
  668. "Unregistered inline_handler %s of %s for %s",
  669. name,
  670. instance.__class__.__name__,
  671. purpose,
  672. )
  673. for name, func in instance.hikka_callback_handlers.copy().items():
  674. if name.lower() in self.callback_handlers and (
  675. hasattr(func, "__self__")
  676. and hasattr(self.callback_handlers[name], "__self__")
  677. and func.__self__.__class__.__name__
  678. == self.callback_handlers[name].__self__.__class__.__name__
  679. ):
  680. del self.callback_handlers[name.lower()]
  681. logger.debug(
  682. "Unregistered callback_handler %s of %s for %s",
  683. name,
  684. instance.__class__.__name__,
  685. purpose,
  686. )
  687. def register_watchers(self, instance: Module):
  688. """Register watcher from instance"""
  689. with contextlib.suppress(AttributeError):
  690. _hikka_client_id_logging_tag = copy.copy(self.client.tg_id) # noqa: F841
  691. for _watcher in self.watchers:
  692. if _watcher.__self__.__class__.__name__ == instance.__class__.__name__:
  693. logger.debug("Removing watcher %s for update", _watcher)
  694. self.watchers.remove(_watcher)
  695. for _watcher in instance.hikka_watchers.values():
  696. self.watchers += [_watcher]
  697. def lookup(
  698. self,
  699. modname: str,
  700. include_dragon: bool = False,
  701. ) -> typing.Union[bool, Module, DragonModule, Library]:
  702. return (
  703. next(
  704. (lib for lib in self.libraries if lib.name.lower() == modname.lower()),
  705. False,
  706. )
  707. or next(
  708. (
  709. mod
  710. for mod in self.modules
  711. if mod.__class__.__name__.lower() == modname.lower()
  712. or mod.name.lower() == modname.lower()
  713. ),
  714. False,
  715. )
  716. or (
  717. next(
  718. (
  719. mod
  720. for mod in self.dragon_modules
  721. if mod.name.lower() == modname.lower()
  722. ),
  723. False,
  724. )
  725. if include_dragon
  726. else False
  727. )
  728. )
  729. @property
  730. def get_approved_channel(self):
  731. return self.__approve.pop(0) if self.__approve else None
  732. def get_prefix(self, userbot: typing.Optional[str] = None) -> str:
  733. """Get prefix for specific userbot. Pass `None` to get Hikka prefix"""
  734. if userbot == "dragon":
  735. key = "dragon.prefix"
  736. default = ","
  737. else:
  738. from . import main
  739. key = main.__name__
  740. default = "."
  741. return self._db.get(key, "command_prefix", default)
  742. async def complete_registration(self, instance: Module):
  743. """Complete registration of instance"""
  744. with contextlib.suppress(AttributeError):
  745. _hikka_client_id_logging_tag = copy.copy(self.client.tg_id) # noqa: F841
  746. instance.allmodules = self
  747. instance.internal_init()
  748. for module in self.modules:
  749. if module.__class__.__name__ == instance.__class__.__name__:
  750. if not self._remove_core_protection and module.__origin__.startswith(
  751. "<core"
  752. ):
  753. raise CoreOverwriteError(
  754. module=(
  755. module.__class__.__name__[:-3]
  756. if module.__class__.__name__.endswith("Mod")
  757. else module.__class__.__name__
  758. )
  759. )
  760. logger.debug("Removing module %s for update", module)
  761. await module.on_unload()
  762. self.modules.remove(module)
  763. for _, method in utils.iter_attrs(module):
  764. if isinstance(method, InfiniteLoop):
  765. method.stop()
  766. logger.debug(
  767. "Stopped loop in module %s, method %s",
  768. module,
  769. method,
  770. )
  771. self.modules += [instance]
  772. def find_alias(
  773. self,
  774. alias: str,
  775. include_legacy: bool = False,
  776. ) -> typing.Optional[str]:
  777. if not alias:
  778. return None
  779. for command_name, _command in self.commands.items():
  780. aliases = []
  781. if getattr(_command, "alias", None) and not (
  782. aliases := getattr(_command, "aliases", None)
  783. ):
  784. aliases = [_command.alias]
  785. if not aliases:
  786. continue
  787. if any(
  788. alias.lower() == _alias.lower()
  789. and alias.lower() not in self._core_commands
  790. for _alias in aliases
  791. ):
  792. return command_name
  793. if alias in self.aliases and include_legacy:
  794. return self.aliases[alias]
  795. return None
  796. def dispatch(self, _command: str) -> typing.Tuple[str, typing.Optional[str]]:
  797. """Dispatch command to appropriate module"""
  798. return next(
  799. (
  800. (cmd, self.commands[cmd.lower()])
  801. for cmd in [
  802. _command,
  803. self.aliases.get(_command.lower()),
  804. self.find_alias(_command),
  805. ]
  806. if cmd and cmd.lower() in self.commands
  807. ),
  808. (_command, None),
  809. )
  810. def send_config(self, skip_hook: bool = False):
  811. """Configure modules"""
  812. for mod in self.modules:
  813. self.send_config_one(mod, skip_hook)
  814. def send_config_one(self, mod: Module, skip_hook: bool = False):
  815. """Send config to single instance"""
  816. with contextlib.suppress(AttributeError):
  817. _hikka_client_id_logging_tag = copy.copy(self.client.tg_id) # noqa: F841
  818. if hasattr(mod, "config"):
  819. modcfg = self._db.get(
  820. mod.__class__.__name__,
  821. "__config__",
  822. {},
  823. )
  824. try:
  825. for conf in mod.config:
  826. with contextlib.suppress(validators.ValidationError):
  827. mod.config.set_no_raise(
  828. conf,
  829. (
  830. modcfg[conf]
  831. if conf in modcfg
  832. else os.environ.get(f"{mod.__class__.__name__}.{conf}")
  833. or mod.config.getdef(conf)
  834. ),
  835. )
  836. except AttributeError:
  837. logger.warning(
  838. "Got invalid config instance. Expected `ModuleConfig`, got %s, %s",
  839. type(mod.config),
  840. mod.config,
  841. )
  842. if not hasattr(mod, "name"):
  843. mod.name = mod.strings["name"]
  844. if skip_hook:
  845. return
  846. if not hasattr(mod, "strings"):
  847. mod.strings = {}
  848. mod.strings = Strings(mod, self.translator)
  849. mod.translator = self.translator
  850. try:
  851. mod.config_complete()
  852. except Exception as e:
  853. logger.exception("Failed to send mod config complete signal due to %s", e)
  854. raise
  855. async def send_ready_one_wrapper(self, *args, **kwargs):
  856. """Wrapper for send_ready_one"""
  857. try:
  858. await self.send_ready_one(*args, **kwargs)
  859. except Exception as e:
  860. logger.exception("Failed to send mod init complete signal due to %s", e)
  861. async def send_ready(self):
  862. """Send all data to all modules"""
  863. await self.inline.register_manager()
  864. await asyncio.gather(
  865. *[self.send_ready_one_wrapper(mod) for mod in self.modules]
  866. )
  867. async def send_ready_one(
  868. self,
  869. mod: Module,
  870. no_self_unload: bool = False,
  871. from_dlmod: bool = False,
  872. ):
  873. with contextlib.suppress(AttributeError):
  874. _hikka_client_id_logging_tag = copy.copy(self.client.tg_id) # noqa: F841
  875. if from_dlmod:
  876. try:
  877. if len(inspect.signature(mod.on_dlmod).parameters) == 2:
  878. await mod.on_dlmod(self.client, self._db)
  879. else:
  880. await mod.on_dlmod()
  881. except Exception:
  882. logger.info("Can't process `on_dlmod` hook", exc_info=True)
  883. try:
  884. if len(inspect.signature(mod.client_ready).parameters) == 2:
  885. await mod.client_ready(self.client, self._db)
  886. else:
  887. await mod.client_ready()
  888. except SelfUnload as e:
  889. if no_self_unload:
  890. raise e
  891. logger.debug("Unloading %s, because it raised SelfUnload", mod)
  892. self.modules.remove(mod)
  893. except SelfSuspend as e:
  894. if no_self_unload:
  895. raise e
  896. logger.debug("Suspending %s, because it raised SelfSuspend", mod)
  897. return
  898. except Exception as e:
  899. logger.exception(
  900. (
  901. "Failed to send mod init complete signal for %s due to %s,"
  902. " attempting unload"
  903. ),
  904. mod,
  905. e,
  906. )
  907. self.modules.remove(mod)
  908. raise
  909. for _, method in utils.iter_attrs(mod):
  910. if isinstance(method, InfiniteLoop):
  911. setattr(method, "module_instance", mod)
  912. if method.autostart:
  913. method.start()
  914. logger.debug("Added module %s to method %s", mod, method)
  915. self.unregister_commands(mod, "update")
  916. self.unregister_raw_handlers(mod, "update")
  917. self.register_commands(mod)
  918. self.register_watchers(mod)
  919. self.register_raw_handlers(mod)
  920. def get_classname(self, name: str) -> str:
  921. return next(
  922. (
  923. module.__class__.__module__
  924. for module in reversed(self.modules)
  925. if name in (module.name, module.__class__.__module__)
  926. ),
  927. name,
  928. )
  929. async def unload_module(self, classname: str) -> typing.List[str]:
  930. """Remove module and all stuff from it"""
  931. worked = []
  932. with contextlib.suppress(AttributeError):
  933. _hikka_client_id_logging_tag = copy.copy(self.client.tg_id) # noqa: F841
  934. for module in self.modules:
  935. if classname.lower() in (
  936. module.name.lower(),
  937. module.__class__.__name__.lower(),
  938. ):
  939. if not self._remove_core_protection and module.__origin__.startswith(
  940. "<core"
  941. ):
  942. raise CoreUnloadError(module.__class__.__name__)
  943. worked += [module.__class__.__name__]
  944. name = module.__class__.__name__
  945. path = os.path.join(
  946. LOADED_MODULES_DIR,
  947. f"{name}_{self.client.tg_id}.py",
  948. )
  949. if os.path.isfile(path):
  950. os.remove(path)
  951. logger.debug("Removed %s file at path %s", name, path)
  952. logger.debug("Removing module %s for unload", module)
  953. self.modules.remove(module)
  954. await module.on_unload()
  955. self.unregister_raw_handlers(module, "unload")
  956. self.unregister_loops(module, "unload")
  957. self.unregister_commands(module, "unload")
  958. self.unregister_watchers(module, "unload")
  959. self.unregister_inline_stuff(module, "unload")
  960. logger.debug("Worked: %s", worked)
  961. return worked
  962. def unregister_loops(self, instance: Module, purpose: str):
  963. for name, method in utils.iter_attrs(instance):
  964. if isinstance(method, InfiniteLoop):
  965. logger.debug(
  966. "Stopping loop for %s in module %s, method %s",
  967. purpose,
  968. instance.__class__.__name__,
  969. name,
  970. )
  971. method.stop()
  972. def unregister_commands(self, instance: Module, purpose: str):
  973. for name, cmd in self.commands.copy().items():
  974. if cmd.__self__.__class__.__name__ == instance.__class__.__name__:
  975. logger.debug(
  976. "Removing command %s of module %s for %s",
  977. name,
  978. instance.__class__.__name__,
  979. purpose,
  980. )
  981. del self.commands[name]
  982. for alias, _command in self.aliases.copy().items():
  983. if _command == name:
  984. del self.aliases[alias]
  985. def unregister_watchers(self, instance: Module, purpose: str):
  986. for _watcher in self.watchers.copy():
  987. if _watcher.__self__.__class__.__name__ == instance.__class__.__name__:
  988. logger.debug(
  989. "Removing watcher %s of module %s for %s",
  990. _watcher,
  991. instance.__class__.__name__,
  992. purpose,
  993. )
  994. self.watchers.remove(_watcher)
  995. def unregister_raw_handlers(self, instance: Module, purpose: str):
  996. """Unregister event handlers for a module"""
  997. for handler in self.client.dispatcher.raw_handlers:
  998. if handler.__self__.__class__.__name__ == instance.__class__.__name__:
  999. self.client.dispatcher.raw_handlers.remove(handler)
  1000. logger.debug(
  1001. "Unregistered raw handler of module %s for %s. ID: %s",
  1002. instance.__class__.__name__,
  1003. purpose,
  1004. handler.id,
  1005. )
  1006. def add_alias(self, alias: str, cmd: str) -> bool:
  1007. """Make an alias"""
  1008. if cmd not in self.commands:
  1009. return False
  1010. self.aliases[alias.lower().strip()] = cmd
  1011. return True
  1012. def remove_alias(self, alias: str) -> bool:
  1013. """Remove an alias"""
  1014. return bool(self.aliases.pop(alias.lower().strip(), None))
  1015. async def log(self, *args, **kwargs):
  1016. """Unnecessary placeholder for logging"""
  1017. async def reload_translations(self) -> bool:
  1018. if not await self.translator.init():
  1019. return False
  1020. for module in self.modules:
  1021. try:
  1022. module.config_complete(reload_dynamic_translate=True)
  1023. except Exception as e:
  1024. logger.debug(
  1025. "Can't complete dynamic translations reload of %s due to %s",
  1026. module,
  1027. e,
  1028. )
  1029. return True