diff --git a/ubot/command.py b/ubot/command.py index 0704205..4a00ba4 100644 --- a/ubot/command.py +++ b/ubot/command.py @@ -14,9 +14,10 @@ class Command: self.extra = args.get("extra", None) self.help = args.get("help", None) or func.__doc__ self.hide_help = args.get("hide_help", False) + self.moderation = args.get("moderation", False) self.owner = args.get("owner", False) self.sudo = args.get("sudo", False) - self.admin = args.get("admin", False) + self.admin = args.get("admin", False) or self.moderation self.nsfw = args.get("nsfw", False) self.nsfw_warning = args.get("nsfw_warning", None) self.pass_nsfw = args.get("pass_nsfw", False) diff --git a/ubot/command_handler.py b/ubot/command_handler.py index 6df6ea4..807d334 100644 --- a/ubot/command_handler.py +++ b/ubot/command_handler.py @@ -57,6 +57,9 @@ class CommandHandler(): pattern_match = search(self.pattern_template.format(f"({'|'.join([escape(i) for i in prefix_list])})", command.pattern + command.pattern_extra, self.micro_bot.me.username), event.raw_text, IGNORECASE|DOTALL) if pattern_match: + if command.moderation and not chat_db.modmode_enabled: + continue + if not (priv_resp := await self.check_privs(event, command, chat_db))[0]: if priv_resp[1]: await event.reply(priv_resp[1]) diff --git a/ubot/database.py b/ubot/database.py index bb50f3f..f6785d4 100644 --- a/ubot/database.py +++ b/ubot/database.py @@ -68,7 +68,7 @@ class ChatWrapper(): # modmode command functions @property - def modmode_enabled(self) -> str: + def modmode_enabled(self) -> bool: return self.chat.modmode_enabled @modmode_enabled.setter diff --git a/ubot/fixes/utils.py b/ubot/fixes/utils.py new file mode 100644 index 0000000..c0b7b16 --- /dev/null +++ b/ubot/fixes/utils.py @@ -0,0 +1,40 @@ +from telethon.tl.types import Channel, Chat + + +async def get_user(event, allow_channel=False): + if event.args: + if event.args.isnumeric(): + user_id = int(event.args) + else: + user_id = None + + try: + user = await event.client.get_entity(user_id or event.args) + + if isinstance(user, Chat) or (isinstance(user, Channel) and not allow_channel): + raise TypeError + + return user + except (ValueError, TypeError): + await event.reply("The ID or username you provided was invalid!") + return + elif event.is_reply: + reply = await event.get_reply_message() + + if reply and reply.sender_id: + try: + user = await event.client.get_entity(reply.sender_id) + + if isinstance(user, Chat) or (isinstance(user, Channel) and not allow_channel): + raise TypeError + + return user + except (ValueError, TypeError): + await event.reply("There was an error getting the user's ID!") + return + else: + await event.reply("Give me a user ID, username or reply!") + return + else: + await event.reply("Give me a user ID, username or reply!") + return diff --git a/ubot/modules/_sudo.py b/ubot/modules/_sudo.py index efa897d..7c66347 100644 --- a/ubot/modules/_sudo.py +++ b/ubot/modules/_sudo.py @@ -11,9 +11,9 @@ from traceback import print_exc import git import psutil from telethon import version -from telethon.tl.types import Channel, Chat from ubot import ldr, startup_time +from ubot.fixes.utils import get_user @ldr.add("eval", owner=True, hide_help=True) @@ -237,43 +237,3 @@ async def rem_sudo(event): async def show_sudo(event): sudo_string = "\n".join([str(user_id) for user_id in ldr.db.sudo_users]) await event.reply(f"**Sudo users:**\n\n{sudo_string}") - - -async def get_user(event): - if event.args: - try: - event.args = int(event.args) - except: - pass - - try: - user = await event.client.get_entity(event.args) - - if isinstance(user, (Chat, Channel)): - raise TypeError - - return user - except (ValueError, TypeError): - await event.reply("The ID or username you provided was invalid!") - return - elif event.is_reply: - reply = await event.get_reply_message() - reply_id = reply.sender_id - - if reply_id: - try: - user = await event.client.get_entity(reply_id) - - if isinstance(user, (Chat, Channel)): - raise TypeError - - return user - except (ValueError, TypeError): - await event.reply("There was an error getting the user's ID!") - return - else: - await event.reply("Sudoing failed!") - return - else: - await event.reply("Give me a user ID, username or reply!") - return