From 4fdc814190454e7f91f006d95e71b89f3686a156 Mon Sep 17 00:00:00 2001 From: Krishna-singhal <65902764+Krishna-Singhal@users.noreply.github.com> Date: Mon, 11 Jan 2021 01:45:58 +0530 Subject: [PATCH] Userbot disable feature and closes #213 (#232) * added feature to disable UserBot in specific chats * userbot disable feature #213 added feature to disable userbot in specific chats * fixed pep8 * -_- * antipattern * fixed some issues * Update system.py * my typos Co-authored-by: rking32 <52490534+rking32@users.noreply.github.com> --- userge/config.py | 2 + .../core/methods/decorators/raw_decorator.py | 4 + userge/plugins/tools/system.py | 102 ++++++++++++++++++ 3 files changed, 108 insertions(+) diff --git a/userge/config.py b/userge/config.py index e9d410fc6..c49c09f8d 100644 --- a/userge/config.py +++ b/userge/config.py @@ -78,6 +78,8 @@ class Config: USE_USER_FOR_CLIENT_CHECKS = False SUDO_ENABLED = False SUDO_USERS: Set[int] = set() + DISABLED_ALL = False + DISABLED_CHATS: Set[int] = set() ALLOWED_COMMANDS: Set[str] = set() ANTISPAM_SENTRY = False RUN_DYNO_SAVER = False diff --git a/userge/core/methods/decorators/raw_decorator.py b/userge/core/methods/decorators/raw_decorator.py index bfdb6f02d..5b062b11a 100644 --- a/userge/core/methods/decorators/raw_decorator.py +++ b/userge/core/methods/decorators/raw_decorator.py @@ -235,6 +235,10 @@ def _build_decorator(self, def decorator(func: _PYROFUNC) -> _PYROFUNC: async def template(r_c: Union['_client.Userge', '_client._UsergeBot'], r_m: RawMessage) -> None: + if Config.DISABLED_ALL and r_m.chat.id != Config.LOG_CHANNEL_ID: + return + if r_m.chat and r_m.chat.id in Config.DISABLED_CHATS: + return await _init(r_c, r_m) _raise = partial(_raise_func, r_c, r_m.chat.id, r_m.message_id) if r_m.chat and r_m.chat.type not in flt.scope: diff --git a/userge/plugins/tools/system.py b/userge/plugins/tools/system.py index 9597fb673..37e30aa4c 100644 --- a/userge/plugins/tools/system.py +++ b/userge/plugins/tools/system.py @@ -18,6 +18,8 @@ from userge.utils import terminate SAVED_SETTINGS = get_collection("CONFIGS") +DISABLED_CHATS = get_collection("DISABLED_CHATS") + MAX_IDLE_TIME = 300 LOG = userge.getLogger(__name__) CHANNEL = userge.getCLogger(__name__) @@ -29,6 +31,14 @@ async def _init() -> None: if d_s: Config.RUN_DYNO_SAVER = bool(d_s['on']) MAX_IDLE_TIME = int(d_s['timeout']) + disabled_all = await SAVED_SETTINGS.find_one({'_id': 'DISABLE_ALL_CHATS'}) + if disabled_all: + Config.DISABLED_ALL = bool(disabled_all['on']) + else: + async for i in DISABLED_CHATS.find(): + if i['_id'] == Config.LOG_CHANNEL_ID: + continue + Config.DISABLED_CHATS.add(i['_id']) @userge.on_cmd('restart', about={ @@ -179,6 +189,98 @@ async def getvar_(message: Message) -> None: await message.edit(f"`var {var_name} forwarded to log channel !`", del_in=3) +@userge.on_cmd("enhere", about={ + 'header': "enable userbot in disabled chat.", + 'flags': {'-all': "Enable Userbot in all chats."}, + 'usage': "{tr}enhere [chat_id | username]\n{tr}enhere -all"}) +async def enable_userbot(message: Message): + if message.flags: + if '-all' in message.flags: + Config.DISABLED_ALL = False + Config.DISABLED_CHATS.clear() + await asyncio.gather( + DISABLED_CHATS.drop(), + SAVED_SETTINGS.update_one( + {'_id': 'DISABLE_ALL_CHATS'}, {"$set": {'on': False}}, upsert=True + ), + message.edit("**Enabled** all chats!", del_in=5)) + else: + await message.err("invalid flag!") + elif message.input_str: + try: + chat = await message.client.get_chat(message.input_str) + except Exception as err: + await message.err(str(err)) + return + if chat.id not in Config.DISABLED_CHATS: + await message.edit("this chat is already enabled!") + else: + Config.DISABLED_CHATS.remove(chat.id) + await asyncio.gather( + DISABLED_CHATS.delete_one( + {'_id': chat.id} + ), + message.edit( + f"CHAT : `{chat.title}` removed from **DISABLED_CHATS**!", + del_in=5, + log=__name__ + ) + ) + else: + await message.err("chat_id not found!") + + +@userge.on_cmd("dishere", about={ + 'header': "enable userbot in current chat.", + 'flags': {'-all': "disable Userbot in all chats."}, + 'usage': "{tr}dishere\n{tr}dishere [chat_id | username]\n{tr}dishere -all"}) +async def disable_userbot(message: Message): + if message.flags: + if '-all' in message.flags: + Config.DISABLED_ALL = True + await asyncio.gather( + SAVED_SETTINGS.update_one( + {'_id': 'DISABLE_ALL_CHATS'}, {"$set": {'on': True}}, upsert=True + ), + message.edit("**Disabled** all chats!", del_in=5)) + else: + await message.err("invalid flag!") + else: + chat = message.chat + if message.input_str: + try: + chat = await message.client.get_chat(message.input_str) + except Exception as err: + await message.err(str(err)) + return + if chat.id in Config.DISABLED_CHATS: + await message.edit("this chat is already disabled!") + elif chat.id == Config.LOG_CHANNEL_ID: + await message.err("can't disabled log channel") + else: + Config.DISABLED_CHATS.add(chat.id) + await asyncio.gather( + DISABLED_CHATS.insert_one({'_id': chat.id, 'title': chat.title}), + message.edit( + f"CHAT : `{chat.title}` added to **DISABLED_CHATS**!", del_in=5, log=__name__ + ) + ) + + +@userge.on_cmd("listdisabled", about={'header': "List all disabled chats."}) +async def view_disabled_chats_(message: Message): + if not Config.DISABLED_CHATS: + await message.edit("**DISABLED_CHATS** not found!", del_in=5) + elif Config.DISABLED_ALL: + # bot will not print this, but dont worry except log channel + await message.edit("All chats are disabled!", del_in=5) + else: + out_str = '🚷 **DISABLED_CHATS** 🚷\n\n' + async for chat in DISABLED_CHATS.find(): + out_str += f" 👥 {chat['title']} 🆔 `{chat['_id']}`\n" + await message.edit(out_str, del_in=0) + + @userge.on_cmd("sleep (\\d+)", about={ 'header': "sleep userge :P", 'usage': "{tr}sleep [timeout in seconds]"}, allow_channels=False)