hikka_security.py 42 KB

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