hikka_security.py 42 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243
  1. # ©️ Dan Gazizullin, 2021-2023
  2. # This file is a part of Hikka Userbot
  3. # 🌐 https://github.com/hikariatama/Hikka
  4. # You can redistribute it and/or modify it under the terms of the GNU AGPLv3
  5. # 🔑 https://www.gnu.org/licenses/agpl-3.0.html
  6. import contextlib
  7. import datetime
  8. import time
  9. import typing
  10. from hikkatl.hints import EntityLike
  11. from hikkatl.tl.types import Message, PeerUser, User
  12. from hikkatl.utils import get_display_name
  13. from .. import loader, main, security, utils
  14. from ..inline.types import InlineCall, InlineMessage
  15. from ..security import (
  16. DEFAULT_PERMISSIONS,
  17. EVERYONE,
  18. GROUP_ADMIN,
  19. GROUP_ADMIN_ADD_ADMINS,
  20. GROUP_ADMIN_BAN_USERS,
  21. GROUP_ADMIN_CHANGE_INFO,
  22. GROUP_ADMIN_DELETE_MESSAGES,
  23. GROUP_ADMIN_INVITE_USERS,
  24. GROUP_ADMIN_PIN_MESSAGES,
  25. GROUP_MEMBER,
  26. GROUP_OWNER,
  27. PM,
  28. SecurityGroup,
  29. )
  30. @loader.tds
  31. class HikkaSecurityMod(loader.Module):
  32. """Control security settings"""
  33. strings = {"name": "HikkaSecurity"}
  34. async def client_ready(self):
  35. self._sgroups: typing.Iterable[str, SecurityGroup] = self.pointer(
  36. "sgroups", {}, item_type=SecurityGroup
  37. )
  38. self._reload_sgroups()
  39. def _reload_sgroups(self):
  40. self._client.dispatcher.security.apply_sgroups(self._sgroups.todict())
  41. async def inline__switch_perm(
  42. self,
  43. call: InlineCall,
  44. command: str,
  45. group: str,
  46. level: bool,
  47. is_inline: bool,
  48. ):
  49. cmd = (
  50. self.allmodules.inline_handlers[command]
  51. if is_inline
  52. else self.allmodules.commands[command]
  53. )
  54. mask = self._db.get(security.__name__, "masks", {}).get(
  55. f"{cmd.__module__}.{cmd.__name__}",
  56. getattr(cmd, "security", security.DEFAULT_PERMISSIONS),
  57. )
  58. bit = security.BITMAP[group.upper()]
  59. if level:
  60. mask |= bit
  61. else:
  62. mask &= ~bit
  63. masks = self._db.get(security.__name__, "masks", {})
  64. masks[f"{cmd.__module__}.{cmd.__name__}"] = mask
  65. self._db.set(security.__name__, "masks", masks)
  66. if (
  67. not self._db.get(security.__name__, "bounding_mask", DEFAULT_PERMISSIONS)
  68. & bit
  69. and level
  70. ):
  71. await call.answer(
  72. (
  73. "Security value set but not applied. Consider enabling this value"
  74. f" in .{'inlinesec' if is_inline else 'security'}"
  75. ),
  76. show_alert=True,
  77. )
  78. else:
  79. await call.answer("Security value set!")
  80. await call.edit(
  81. self.strings("permissions").format(
  82. utils.escape_html(
  83. f"@{self.inline.bot_username} " if is_inline else self.get_prefix()
  84. ),
  85. command,
  86. ),
  87. reply_markup=self._build_markup(cmd, is_inline),
  88. )
  89. async def inline__switch_perm_bm(
  90. self,
  91. call: InlineCall,
  92. group: str,
  93. level: bool,
  94. is_inline: bool,
  95. ):
  96. mask = self._db.get(security.__name__, "bounding_mask", DEFAULT_PERMISSIONS)
  97. bit = security.BITMAP[group.upper()]
  98. if level:
  99. mask |= bit
  100. else:
  101. mask &= ~bit
  102. self._db.set(security.__name__, "bounding_mask", mask)
  103. await call.answer("Bounding mask value set!")
  104. await call.edit(
  105. self.strings("global"),
  106. reply_markup=self._build_markup_global(is_inline),
  107. )
  108. def _build_markup(
  109. self,
  110. command: callable,
  111. is_inline: bool = False,
  112. ) -> typing.List[typing.List[dict]]:
  113. perms = self._get_current_perms(command, is_inline)
  114. return (
  115. utils.chunks(
  116. [
  117. {
  118. "text": f"{'✅' if level else '🚫'} {self.strings[group]}",
  119. "callback": self.inline__switch_perm,
  120. "args": (
  121. command.__name__.rsplit("_inline_handler", maxsplit=1)[0],
  122. group,
  123. not level,
  124. is_inline,
  125. ),
  126. }
  127. for group, level in perms.items()
  128. ],
  129. 2,
  130. )
  131. + [[{"text": self.strings("close_menu"), "action": "close"}]]
  132. if is_inline
  133. else utils.chunks(
  134. [
  135. {
  136. "text": f"{'✅' if level else '🚫'} {self.strings[group]}",
  137. "callback": self.inline__switch_perm,
  138. "args": (
  139. command.__name__.rsplit("cmd", maxsplit=1)[0],
  140. group,
  141. not level,
  142. is_inline,
  143. ),
  144. }
  145. for group, level in perms.items()
  146. ],
  147. 2,
  148. )
  149. + [[{
  150. "text": self.strings("close_menu"),
  151. "action": "close",
  152. }]]
  153. )
  154. def _build_markup_global(
  155. self,
  156. is_inline: bool = False,
  157. ) -> typing.List[typing.List[dict]]:
  158. return utils.chunks(
  159. [
  160. {
  161. "text": f"{'✅' if level else '🚫'} {self.strings[group]}",
  162. "callback": self.inline__switch_perm_bm,
  163. "args": (group, not level, is_inline),
  164. }
  165. for group, level in self._get_current_bm(is_inline).items()
  166. ],
  167. 2,
  168. ) + [[{"text": self.strings("close_menu"), "action": "close"}]]
  169. def _get_current_bm(self, is_inline: bool = False) -> dict:
  170. return self._perms_map(
  171. self._db.get(security.__name__, "bounding_mask", DEFAULT_PERMISSIONS),
  172. is_inline,
  173. )
  174. @staticmethod
  175. def _perms_map(perms: int, is_inline: bool) -> dict:
  176. return (
  177. {"everyone": bool(perms & EVERYONE)}
  178. if is_inline
  179. else {
  180. "group_owner": bool(perms & GROUP_OWNER),
  181. "group_admin_add_admins": bool(perms & GROUP_ADMIN_ADD_ADMINS),
  182. "group_admin_change_info": bool(perms & GROUP_ADMIN_CHANGE_INFO),
  183. "group_admin_ban_users": bool(perms & GROUP_ADMIN_BAN_USERS),
  184. "group_admin_delete_messages": bool(
  185. perms & GROUP_ADMIN_DELETE_MESSAGES
  186. ),
  187. "group_admin_pin_messages": bool(perms & GROUP_ADMIN_PIN_MESSAGES),
  188. "group_admin_invite_users": bool(perms & GROUP_ADMIN_INVITE_USERS),
  189. "group_admin": bool(perms & GROUP_ADMIN),
  190. "group_member": bool(perms & GROUP_MEMBER),
  191. "pm": bool(perms & PM),
  192. "everyone": bool(perms & EVERYONE),
  193. }
  194. )
  195. def _get_current_perms(
  196. self,
  197. command: callable,
  198. is_inline: bool = False,
  199. ) -> dict:
  200. return self._perms_map(
  201. self._db.get(security.__name__, "masks", {}).get(
  202. f"{command.__module__}.{command.__name__}",
  203. getattr(command, "security", self._client.dispatcher.security.default),
  204. ),
  205. is_inline,
  206. )
  207. @loader.command()
  208. async def newsgroup(self, message: Message):
  209. if not (args := utils.get_args_raw(message)):
  210. await utils.answer(message, self.strings("no_args"))
  211. return
  212. if not all([i.isalnum() for i in args]):
  213. await utils.answer(message, self.strings("invalid_name"))
  214. return
  215. if args in self._sgroups:
  216. await utils.answer(
  217. message, self.strings("sgroup_already_exists").format(args)
  218. )
  219. return
  220. self._sgroups[args] = SecurityGroup(args, [], [])
  221. self._reload_sgroups()
  222. await utils.answer(message, self.strings("created_sgroup").format(args))
  223. @loader.command()
  224. async def sgroups(self, message: Message):
  225. await utils.answer(
  226. message,
  227. self.strings("sgroups_list").format(
  228. "\n".join(
  229. self.strings("sgroup_li").format(
  230. group.name, len(group.users), len(group.permissions)
  231. )
  232. for group in self._sgroups.values()
  233. )
  234. ),
  235. )
  236. @loader.command()
  237. async def sgroup(self, message: Message):
  238. if not (args := utils.get_args_raw(message)):
  239. await utils.answer(message, self.strings("no_args"))
  240. return
  241. if not (group := self._sgroups.get(args)):
  242. await utils.answer(message, self.strings("sgroup_not_found").format(args))
  243. return
  244. await utils.answer(
  245. message,
  246. self.strings("sgroup_info").format(
  247. group.name,
  248. (
  249. self.strings("users_list").format(
  250. "\n".join([
  251. self.strings("li").format(
  252. utils.get_entity_url(
  253. await self._client.get_entity(user, exp=0)
  254. ),
  255. utils.escape_html(
  256. get_display_name(
  257. await self._client.get_entity(user, exp=0)
  258. )
  259. ),
  260. )
  261. for user in group.users
  262. ])
  263. )
  264. if group.users
  265. else self.strings("no_users")
  266. ),
  267. (
  268. self.strings("permissions_list").format(
  269. "\n".join(
  270. "<emoji document_id=4974307891025543730>▫️</emoji>"
  271. " <b>{}</b> <code>{}</code> <b>{}</b>".format(
  272. self.strings(rule["rule_type"]),
  273. rule["rule"],
  274. (
  275. (
  276. self.strings("until")
  277. + " "
  278. + self._convert_time_abs(rule["expires"])
  279. )
  280. if rule["expires"]
  281. else self.strings("forever")
  282. ),
  283. )
  284. for rule in group.permissions
  285. )
  286. )
  287. if group.permissions
  288. else self.strings("no_permissions")
  289. ),
  290. ),
  291. )
  292. @loader.command()
  293. async def delsgroup(self, message: Message):
  294. if not (args := utils.get_args_raw(message)):
  295. await utils.answer(message, self.strings("no_args"))
  296. return
  297. if not self._sgroups.get(args):
  298. await utils.answer(message, self.strings("sgroup_not_found").format(args))
  299. return
  300. del self._sgroups[args]
  301. self._reload_sgroups()
  302. await utils.answer(message, self.strings("deleted_sgroup").format(args))
  303. @loader.command()
  304. async def sgroupadd(self, message: Message):
  305. if not (args := utils.get_args_raw(message)):
  306. await utils.answer(message, self.strings("no_args"))
  307. return
  308. if len(args.split()) >= 2:
  309. group, user = args.split()
  310. if user.isdigit():
  311. user = int(user)
  312. try:
  313. user = await self._client.get_entity(user, exp=0)
  314. except ValueError:
  315. await utils.answer(message, self.strings("no_args"))
  316. return
  317. else:
  318. if not message.is_reply:
  319. await utils.answer(message, self.strings("no_args"))
  320. return
  321. group, user = args, await (await message.get_reply_message()).get_sender()
  322. if not (group := self._sgroups.get(group)):
  323. await utils.answer(message, self.strings("sgroup_not_found").format(group))
  324. return
  325. if user.id in group.users:
  326. await utils.answer(
  327. message,
  328. self.strings("user_already_in_sgroup").format(
  329. utils.escape_html(get_display_name(user)),
  330. group.name,
  331. ),
  332. )
  333. return
  334. group.users.append(user.id)
  335. self._sgroups[group.name] = group
  336. self._reload_sgroups()
  337. await utils.answer(
  338. message,
  339. self.strings("user_added_to_sgroup").format(
  340. utils.escape_html(get_display_name(user)),
  341. group.name,
  342. ),
  343. )
  344. @loader.command()
  345. async def sgroupdel(self, message: Message):
  346. if not (args := utils.get_args_raw(message)):
  347. await utils.answer(message, self.strings("no_args"))
  348. return
  349. if len(args.split()) >= 2:
  350. group, user = args.split()
  351. if user.isdigit():
  352. user = int(user)
  353. try:
  354. user = await self._client.get_entity(user, exp=0)
  355. except ValueError:
  356. await utils.answer(message, self.strings("no_args"))
  357. return
  358. else:
  359. if not message.is_reply:
  360. await utils.answer(message, self.strings("no_args"))
  361. return
  362. group, user = args, await (await message.get_reply_message()).get_sender()
  363. if not (group := self._sgroups.get(group)):
  364. await utils.answer(message, self.strings("sgroup_not_found").format(group))
  365. return
  366. if user.id not in group.users:
  367. await utils.answer(
  368. message,
  369. self.strings("user_not_in_sgroup").format(
  370. utils.escape_html(get_display_name(user)),
  371. group.name,
  372. ),
  373. )
  374. return
  375. group.users.remove(user.id)
  376. self._sgroups[group.name] = group
  377. self._reload_sgroups()
  378. await utils.answer(
  379. message,
  380. self.strings("user_removed_from_sgroup").format(
  381. utils.escape_html(get_display_name(user)),
  382. group.name,
  383. ),
  384. )
  385. @loader.command()
  386. async def security(self, message: Message):
  387. if (
  388. args := utils.get_args_raw(message).lower().strip()
  389. ) and args not in self.allmodules.commands:
  390. await utils.answer(message, self.strings("no_command").format(args))
  391. return
  392. if not args:
  393. await self.inline.form(
  394. self.strings("global"),
  395. reply_markup=self._build_markup_global(),
  396. message=message,
  397. ttl=5 * 60,
  398. )
  399. return
  400. await self.inline.form(
  401. self.strings("permissions").format(self.get_prefix(), args),
  402. reply_markup=self._build_markup(self.allmodules.commands[args]),
  403. message=message,
  404. ttl=5 * 60,
  405. )
  406. @loader.command()
  407. async def inlinesec(self, message: Message):
  408. if not (args := utils.get_args_raw(message).lower().strip()):
  409. await self.inline.form(
  410. self.strings("global"),
  411. reply_markup=self._build_markup_global(True),
  412. message=message,
  413. ttl=5 * 60,
  414. )
  415. return
  416. if args not in self.allmodules.inline_handlers:
  417. await utils.answer(message, self.strings("no_command").format(args))
  418. return
  419. await self.inline.form(
  420. self.strings("permissions").format(f"@{self.inline.bot_username} ", args),
  421. reply_markup=self._build_markup(
  422. self.allmodules.inline_handlers[args],
  423. True,
  424. ),
  425. message=message,
  426. ttl=5 * 60,
  427. )
  428. async def _resolve_user(self, message: Message):
  429. if not (args := utils.get_args_raw(message)) and not (
  430. reply := await message.get_reply_message()
  431. ):
  432. await utils.answer(message, self.strings("no_user"))
  433. return
  434. user = None
  435. if args:
  436. with contextlib.suppress(Exception):
  437. if str(args).isdigit():
  438. args = int(args)
  439. user = await self._client.get_entity(args, exp=0)
  440. if user is None:
  441. try:
  442. user = await self._client.get_entity(reply.sender_id, exp=0)
  443. except ValueError:
  444. user = await reply.get_sender()
  445. if not isinstance(user, (User, PeerUser)):
  446. await utils.answer(message, self.strings("not_a_user"))
  447. return
  448. if user.id == self.tg_id:
  449. await utils.answer(message, self.strings("self"))
  450. return
  451. return user
  452. async def _add_to_group(
  453. self,
  454. message: typing.Union[Message, InlineCall],
  455. group: str,
  456. confirmed: bool = False,
  457. user: int = None,
  458. ):
  459. if user is None and not (user := await self._resolve_user(message)):
  460. return
  461. if isinstance(user, int):
  462. user = await self._client.get_entity(user, exp=0)
  463. if not confirmed:
  464. await self.inline.form(
  465. self.strings("warning").format(
  466. user.id,
  467. utils.escape_html(get_display_name(user)),
  468. group,
  469. ),
  470. message=message,
  471. ttl=10 * 60,
  472. reply_markup=[
  473. {
  474. "text": self.strings("cancel"),
  475. "action": "close",
  476. },
  477. {
  478. "text": self.strings("confirm"),
  479. "callback": self._add_to_group,
  480. "args": (group, True, user.id),
  481. },
  482. ],
  483. )
  484. return
  485. if user.id not in getattr(self._client.dispatcher.security, group):
  486. getattr(self._client.dispatcher.security, group).append(user.id)
  487. await message.edit(
  488. (
  489. self.strings(f"{group}_added").format(
  490. user.id,
  491. utils.escape_html(get_display_name(user)),
  492. )
  493. + "\n\n"
  494. + self.strings("suggest_nonick")
  495. ),
  496. reply_markup=[
  497. {
  498. "text": self.strings("cancel"),
  499. "action": "close",
  500. },
  501. {
  502. "text": self.strings("enable_nonick_btn"),
  503. "callback": self._enable_nonick,
  504. "args": (user,),
  505. },
  506. ],
  507. )
  508. async def _enable_nonick(self, call: InlineCall, user: User):
  509. self._db.set(
  510. main.__name__,
  511. "nonickusers",
  512. list(set(self._db.get(main.__name__, "nonickusers", []) + [user.id])),
  513. )
  514. await call.edit(
  515. self.strings("user_nn").format(
  516. user.id,
  517. utils.escape_html(get_display_name(user)),
  518. )
  519. )
  520. await call.unload()
  521. @loader.command()
  522. async def owneradd(self, message: Message):
  523. await self._add_to_group(message, "owner")
  524. @loader.command()
  525. async def ownerrm(self, message: Message):
  526. if not (user := await self._resolve_user(message)):
  527. return
  528. if user.id in self._client.dispatcher.security.owner:
  529. self._client.dispatcher.security.owner.remove(user.id)
  530. await utils.answer(
  531. message,
  532. self.strings("owner_removed").format(
  533. user.id,
  534. utils.escape_html(get_display_name(user)),
  535. ),
  536. )
  537. @loader.command()
  538. async def ownerlist(self, message: Message):
  539. _resolved_users = []
  540. for user in set(self._client.dispatcher.security.owner + [self.tg_id]):
  541. with contextlib.suppress(Exception):
  542. _resolved_users += [await self._client.get_entity(user, exp=0)]
  543. if not _resolved_users:
  544. await utils.answer(message, self.strings("no_owner"))
  545. return
  546. await utils.answer(
  547. message,
  548. self.strings("owner_list").format(
  549. "\n".join([
  550. self.strings("li").format(
  551. i.id, utils.escape_html(get_display_name(i))
  552. )
  553. for i in _resolved_users
  554. ])
  555. ),
  556. )
  557. def _lookup(self, needle: str) -> str:
  558. return (
  559. (
  560. []
  561. if needle.lower().startswith(self.get_prefix())
  562. else (
  563. [f"module/{self.lookup(needle).__class__.__name__}"]
  564. if self.lookup(needle)
  565. else []
  566. )
  567. )
  568. + (
  569. [f"command/{needle.lower().strip(self.get_prefix())}"]
  570. if needle.lower().strip(self.get_prefix()) in self.allmodules.commands
  571. else []
  572. )
  573. + (
  574. [f"inline/{needle.lower().strip('@')}"]
  575. if needle.lower().strip("@") in self.allmodules.inline_handlers
  576. else []
  577. )
  578. )
  579. @staticmethod
  580. def _extract_time(args: list) -> int:
  581. for suffix, quantifier in [
  582. ("d", 24 * 60 * 60),
  583. ("h", 60 * 60),
  584. ("m", 60),
  585. ("s", 1),
  586. ]:
  587. duration = next(
  588. (
  589. int(arg.rsplit(suffix, maxsplit=1)[0])
  590. for arg in args
  591. if arg.endswith(suffix)
  592. and arg.rsplit(suffix, maxsplit=1)[0].isdigit()
  593. ),
  594. None,
  595. )
  596. if duration is not None:
  597. return duration * quantifier
  598. return 0
  599. def _convert_time_abs(self, timestamp: int) -> str:
  600. return (
  601. self.strings("forever")
  602. if not timestamp
  603. else datetime.datetime.fromtimestamp(timestamp).strftime(
  604. "%Y-%m-%d %H:%M:%S"
  605. )
  606. )
  607. def _convert_time(self, duration: int) -> str:
  608. return (
  609. self.strings("forever")
  610. if not duration or duration < 0
  611. else (
  612. (
  613. f"{duration // (24 * 60 * 60)} "
  614. + self.strings(
  615. f"day{'s' if duration // (24 * 60 * 60) > 1 else ''}"
  616. )
  617. )
  618. if duration >= 24 * 60 * 60
  619. else (
  620. (
  621. f"{duration // (60 * 60)} "
  622. + self.strings(
  623. f"hour{'s' if duration // (60 * 60) > 1 else ''}"
  624. )
  625. )
  626. if duration >= 60 * 60
  627. else (
  628. (
  629. f"{duration // 60} "
  630. + self.strings(f"minute{'s' if duration // 60 > 1 else ''}")
  631. )
  632. if duration >= 60
  633. else (
  634. f"{duration} "
  635. + self.strings(f"second{'s' if duration > 1 else ''}")
  636. )
  637. )
  638. )
  639. )
  640. )
  641. async def _add_rule(
  642. self,
  643. call: InlineCall,
  644. target_type: str,
  645. target: typing.Union[EntityLike, str],
  646. rule: str,
  647. duration: int,
  648. ):
  649. if rule.startswith("inline") and target_type == "chat":
  650. await call.edit(self.strings("chat_inline"))
  651. return
  652. if target_type == "sgroup":
  653. group = self._sgroups[target]
  654. group.permissions.append({
  655. "target": target,
  656. "rule_type": rule.split("/")[0],
  657. "rule": rule.split("/", maxsplit=1)[1],
  658. "expires": int(time.time() + duration) if duration else 0,
  659. "entity_name": group.name,
  660. "entity_url": "",
  661. })
  662. self._reload_sgroups()
  663. else:
  664. self._client.dispatcher.security.add_rule(
  665. target_type,
  666. target,
  667. rule,
  668. duration,
  669. )
  670. await call.edit(
  671. self.strings("rule_added").format(
  672. self.strings(target_type),
  673. utils.get_entity_url(target) if not isinstance(target, str) else "",
  674. utils.escape_html(
  675. get_display_name(target) if not isinstance(target, str) else target
  676. ),
  677. self.strings(rule.split("/", maxsplit=1)[0]),
  678. (
  679. (
  680. f"@{self.inline.bot_username} "
  681. if rule.split("/", maxsplit=1)[0] == "inline"
  682. else ""
  683. )
  684. + rule.split("/", maxsplit=1)[1]
  685. ),
  686. (
  687. (self.strings("for") + " " + self._convert_time(duration))
  688. if duration
  689. else self.strings("forever")
  690. ),
  691. )
  692. )
  693. async def _confirm(
  694. self,
  695. obj: typing.Union[Message, InlineMessage],
  696. target_type: str,
  697. target: typing.Union[EntityLike, str],
  698. rule: str,
  699. duration: int,
  700. ):
  701. await utils.answer(
  702. obj,
  703. self.strings("confirm_rule").format(
  704. self.strings(target_type),
  705. utils.get_entity_url(target) if not isinstance(target, str) else "",
  706. utils.escape_html(
  707. get_display_name(target) if not isinstance(target, str) else target
  708. ),
  709. self.strings(rule.split("/", maxsplit=1)[0]),
  710. (
  711. (
  712. f"@{self.inline.bot_username} "
  713. if rule.split("/", maxsplit=1)[0] == "inline"
  714. else ""
  715. )
  716. + rule.split("/", maxsplit=1)[1]
  717. ),
  718. (
  719. (self.strings("for") + " " + self._convert_time(duration))
  720. if duration
  721. else self.strings("forever")
  722. ),
  723. ),
  724. reply_markup=[
  725. {
  726. "text": self.strings("confirm_btn"),
  727. "callback": self._add_rule,
  728. "args": (target_type, target, rule, duration),
  729. },
  730. {"text": self.strings("cancel_btn"), "action": "close"},
  731. ],
  732. )
  733. async def _tsec_chat(self, message: Message, args: list):
  734. if len(args) == 1 and message.is_private:
  735. await utils.answer(message, self.strings("no_target"))
  736. return
  737. if len(args) >= 2:
  738. try:
  739. if not args[1].isdigit() and not args[1].startswith("@"):
  740. raise ValueError
  741. target = await self._client.get_entity(
  742. int(args[1]) if args[1].isdigit() else args[1],
  743. exp=0,
  744. )
  745. except (ValueError, TypeError):
  746. if not message.is_private:
  747. target = await self._client.get_entity(message.peer_id, exp=0)
  748. else:
  749. await utils.answer(message, self.strings("no_target"))
  750. return
  751. if not (
  752. possible_rules := utils.array_sum([self._lookup(arg) for arg in args[1:]])
  753. ):
  754. await utils.answer(message, self.strings("no_rule"))
  755. return
  756. duration = self._extract_time(args)
  757. if len(possible_rules) > 1:
  758. await self.inline.form(
  759. message=message,
  760. text=self.strings("multiple_rules").format(
  761. "\n".join(
  762. "🛡 <b>{}</b> <code>{}</code>".format(
  763. self.strings(rule.split("/")[0]).capitalize(),
  764. rule.split("/", maxsplit=1)[1],
  765. )
  766. for rule in possible_rules
  767. )
  768. ),
  769. reply_markup=utils.chunks(
  770. [
  771. {
  772. "text": "🛡 {} {}".format(
  773. self.strings(rule.split("/")[0]).capitalize(),
  774. rule.split("/", maxsplit=1)[1],
  775. ),
  776. "callback": self._confirm,
  777. "args": ("chat", target, rule, duration),
  778. }
  779. for rule in possible_rules
  780. ],
  781. 3,
  782. ),
  783. )
  784. return
  785. await self._confirm(message, "chat", target, possible_rules[0], duration)
  786. async def _tsec_sgroup(self, message: Message, args: list):
  787. if len(args) <= 1:
  788. await utils.answer(message, self.strings("no_target"))
  789. return
  790. if (target := args[1]) not in self._sgroups:
  791. await utils.answer(message, self.strings("sgroup_not_found").format(target))
  792. return
  793. if not (
  794. possible_rules := utils.array_sum([self._lookup(arg) for arg in args[1:]])
  795. ):
  796. await utils.answer(message, self.strings("no_rule"))
  797. return
  798. duration = self._extract_time(args)
  799. if len(possible_rules) > 1:
  800. await self.inline.form(
  801. message=message,
  802. text=self.strings("multiple_rules").format(
  803. "\n".join(
  804. "🛡 <b>{}</b> <code>{}</code>".format(
  805. self.strings(rule.split("/")[0]).capitalize(),
  806. rule.split("/", maxsplit=1)[1],
  807. )
  808. for rule in possible_rules
  809. )
  810. ),
  811. reply_markup=utils.chunks(
  812. [
  813. {
  814. "text": "🛡 {} {}".format(
  815. self.strings(rule.split("/")[0]).capitalize(),
  816. rule.split("/", maxsplit=1)[1],
  817. ),
  818. "callback": self._confirm,
  819. "args": ("sgroup", target, rule, duration),
  820. }
  821. for rule in possible_rules
  822. ],
  823. 3,
  824. ),
  825. )
  826. return
  827. await self._confirm(message, "sgroup", target, possible_rules[0], duration)
  828. async def _tsec_user(self, message: Message, args: list):
  829. if len(args) == 1 and not message.is_private and not message.is_reply:
  830. await utils.answer(message, self.strings("no_target"))
  831. return
  832. if len(args) >= 2:
  833. try:
  834. if not args[1].isdigit() and not args[1].startswith("@"):
  835. raise ValueError
  836. target = await self._client.get_entity(
  837. int(args[1]) if args[1].isdigit() else args[1],
  838. exp=0,
  839. )
  840. except (ValueError, TypeError):
  841. if message.is_private:
  842. target = await self._client.get_entity(message.peer_id, exp=0)
  843. elif message.is_reply:
  844. target = await self._client.get_entity(
  845. (await message.get_reply_message()).sender_id,
  846. exp=0,
  847. )
  848. else:
  849. await utils.answer(message, self.strings("no_target"))
  850. return
  851. if target.id in self._client.dispatcher.security.owner:
  852. await utils.answer(message, self.strings("owner_target"))
  853. return
  854. duration = self._extract_time(args)
  855. if not (
  856. possible_rules := utils.array_sum([self._lookup(arg) for arg in args[1:]])
  857. ):
  858. await utils.answer(message, self.strings("no_rule"))
  859. return
  860. if len(possible_rules) > 1:
  861. await self.inline.form(
  862. message=message,
  863. text=self.strings("multiple_rules").format(
  864. "\n".join(
  865. "🛡 <b>{}</b> <code>{}</code>".format(
  866. self.strings(rule.split("/")[0]).capitalize(),
  867. rule.split("/", maxsplit=1)[1],
  868. )
  869. for rule in possible_rules
  870. )
  871. ),
  872. reply_markup=utils.chunks(
  873. [
  874. {
  875. "text": "🛡 {} {}".format(
  876. self.strings(rule.split("/")[0]).capitalize(),
  877. rule.split("/", maxsplit=1)[1],
  878. ),
  879. "callback": self._confirm,
  880. "args": ("user", target, rule, duration),
  881. }
  882. for rule in possible_rules
  883. ],
  884. 3,
  885. ),
  886. )
  887. return
  888. await self._confirm(message, "user", target, possible_rules[0], duration)
  889. @loader.command()
  890. async def tsecrm(self, message: Message):
  891. if (
  892. not self._client.dispatcher.security.tsec_chat
  893. and not self._client.dispatcher.security.tsec_user
  894. ):
  895. await utils.answer(message, self.strings("no_rules"))
  896. return
  897. if not (args := utils.get_args(message)) or args[0] not in {
  898. "user",
  899. "chat",
  900. "sgroup",
  901. }:
  902. await utils.answer(message, self.strings("no_target"))
  903. return
  904. if args[0] == "user":
  905. if not message.is_private and not message.is_reply:
  906. await utils.answer(message, self.strings("no_target"))
  907. return
  908. if message.is_private:
  909. target = await self._client.get_entity(message.peer_id, exp=0)
  910. elif message.is_reply:
  911. target = await self._client.get_entity(
  912. (await message.get_reply_message()).sender_id,
  913. exp=0,
  914. )
  915. if not self._client.dispatcher.security.remove_rule(
  916. "user",
  917. target.id,
  918. args[1],
  919. ):
  920. await utils.answer(message, self.strings("no_rules"))
  921. return
  922. await utils.answer(
  923. message,
  924. self.strings("rule_removed").format(
  925. utils.get_entity_url(target),
  926. utils.escape_html(get_display_name(target)),
  927. utils.escape_html(args[1]),
  928. ),
  929. )
  930. return
  931. if args[0] == "sgroup":
  932. if len(args) < 3 or args[1] not in self._sgroups:
  933. await utils.answer(message, self.strings("no_target"))
  934. return
  935. group = self._sgroups[args[1]]
  936. _any = False
  937. for rule in group.permissions:
  938. if rule["rule"] == args[2]:
  939. group.permissions.remove(rule)
  940. _any = True
  941. if not _any:
  942. await utils.answer(message, self.strings("no_rules"))
  943. return
  944. self._reload_sgroups()
  945. await utils.answer(
  946. message,
  947. self.strings("rule_removed").format(
  948. "",
  949. utils.escape_html(group.name),
  950. utils.escape_html(args[2]),
  951. ),
  952. )
  953. return
  954. if message.is_private:
  955. await utils.answer(message, self.strings("no_target"))
  956. return
  957. target = await self._client.get_entity(message.peer_id, exp=0)
  958. if not self._client.dispatcher.security.remove_rule("chat", target.id, args[1]):
  959. await utils.answer(message, self.strings("no_rules"))
  960. return
  961. await utils.answer(
  962. message,
  963. self.strings("rule_removed").format(
  964. utils.get_entity_url(target),
  965. utils.escape_html(get_display_name(target)),
  966. utils.escape_html(args[1]),
  967. ),
  968. )
  969. @loader.command()
  970. async def tsecclr(self, message: Message):
  971. if (
  972. not self._client.dispatcher.security.tsec_chat
  973. and not self._client.dispatcher.security.tsec_user
  974. ):
  975. await utils.answer(message, self.strings("no_rules"))
  976. return
  977. if (
  978. not (args := utils.get_args(message))
  979. or not (args := args[0])
  980. or args not in {"user", "chat", "sgroup"}
  981. ):
  982. await utils.answer(message, self.strings("no_target"))
  983. return
  984. if args == "user":
  985. if not message.is_private and not message.is_reply:
  986. await utils.answer(message, self.strings("no_target"))
  987. return
  988. if message.is_private:
  989. target = await self._client.get_entity(message.peer_id, exp=0)
  990. elif message.is_reply:
  991. target = await self._client.get_entity(
  992. (await message.get_reply_message()).sender_id,
  993. exp=0,
  994. )
  995. if not self._client.dispatcher.security.remove_rules("user", target.id):
  996. await utils.answer(message, self.strings("no_rules"))
  997. return
  998. await utils.answer(
  999. message,
  1000. self.strings("rules_removed").format(
  1001. utils.get_entity_url(target),
  1002. utils.escape_html(get_display_name(target)),
  1003. ),
  1004. )
  1005. return
  1006. if args == "sgroup":
  1007. group = utils.get_args(message)[1]
  1008. if not (group := self._sgroups.get(group)):
  1009. await utils.answer(message, self.strings("no_target"))
  1010. return
  1011. group.permissions.clear()
  1012. self._sgroups[group.name] = group
  1013. self._reload_sgroups()
  1014. await utils.answer(
  1015. message,
  1016. self.strings("rules_removed").format(
  1017. "",
  1018. utils.escape_html(group.name),
  1019. ),
  1020. )
  1021. return
  1022. if message.is_private:
  1023. await utils.answer(message, self.strings("no_target"))
  1024. return
  1025. target = await self._client.get_entity(message.peer_id, exp=0)
  1026. if not self._client.dispatcher.security.remove_rules("chat", target.id):
  1027. await utils.answer(message, self.strings("no_rules"))
  1028. return
  1029. await utils.answer(
  1030. message,
  1031. self.strings("rules_removed").format(
  1032. utils.get_entity_url(target),
  1033. utils.escape_html(get_display_name(target)),
  1034. ),
  1035. )
  1036. @loader.command()
  1037. async def tsec(self, message: Message):
  1038. if not (args := utils.get_args(message)):
  1039. if (
  1040. not self._client.dispatcher.security.tsec_chat
  1041. and not self._client.dispatcher.security.tsec_user
  1042. ):
  1043. await utils.answer(message, self.strings("no_rules"))
  1044. return
  1045. await utils.answer(
  1046. message,
  1047. self.strings("rules").format(
  1048. "\n".join(
  1049. [
  1050. "<emoji document_id=6037355667365300960>👥</emoji> <b><a"
  1051. " href='{}'>{}</a> {} {} {}</b> <code>{}</code>".format(
  1052. rule["entity_url"],
  1053. utils.escape_html(rule["entity_name"]),
  1054. self._convert_time(int(rule["expires"] - time.time())),
  1055. self.strings("for"),
  1056. self.strings(rule["rule_type"]),
  1057. rule["rule"],
  1058. )
  1059. for rule in self._client.dispatcher.security.tsec_chat
  1060. ]
  1061. + [
  1062. "<emoji document_id=6037122016849432064>👤</emoji> <b><a"
  1063. " href='{}'>{}</a> {} {} {}</b> <code>{}</code>".format(
  1064. rule["entity_url"],
  1065. utils.escape_html(rule["entity_name"]),
  1066. self._convert_time(int(rule["expires"] - time.time())),
  1067. self.strings("for"),
  1068. self.strings(rule["rule_type"]),
  1069. rule["rule"],
  1070. )
  1071. for rule in self._client.dispatcher.security.tsec_user
  1072. ]
  1073. + [
  1074. "\n".join([
  1075. "<emoji document_id=5870704313440932932>🔒</emoji>"
  1076. " <code>{}</code> <b>{} {} {}</b> <code>{}</code>".format(
  1077. utils.escape_html(group.name),
  1078. self._convert_time(
  1079. int(rule["expires"] - time.time())
  1080. ),
  1081. self.strings("for"),
  1082. self.strings(rule["rule_type"]),
  1083. rule["rule"],
  1084. )
  1085. for rule in group.permissions
  1086. ])
  1087. for group in self._sgroups.values()
  1088. ]
  1089. )
  1090. ),
  1091. )
  1092. return
  1093. if args[0] not in {"user", "chat", "sgroup"}:
  1094. await utils.answer(message, self.strings("what"))
  1095. return
  1096. await getattr(self, f"_tsec_{args[0]}")(message, args)