From 9f26ed107ca6297162974ff085d225047baa99ff Mon Sep 17 00:00:00 2001 From: lzzy12 Date: Mon, 14 Oct 2019 21:53:31 +0530 Subject: [PATCH] Allow owner to unauthorize a user or a chat Signed-off-by: lzzy12 --- bot/__init__.py | 4 +-- bot/helper/telegram_helper/bot_commands.py | 1 + bot/modules/authorize.py | 36 ++++++++++++++++++++-- 3 files changed, 37 insertions(+), 4 deletions(-) diff --git a/bot/__init__.py b/bot/__init__.py index e3914e3..8636ec9 100644 --- a/bot/__init__.py +++ b/bot/__init__.py @@ -49,13 +49,13 @@ status_reply_dict = {} # Value: An object of DownloadStatus download_dict = {} # Stores list of users and chats the bot is authorized to use in -AUTHORIZED_CHATS = [] +AUTHORIZED_CHATS = set() if os.path.exists('authorized_chats.txt'): with open('authorized_chats.txt', 'r+') as f: lines = f.readlines() for line in lines: # LOGGER.info(line.split()) - AUTHORIZED_CHATS.append(int(line.split()[0])) + AUTHORIZED_CHATS.add(int(line.split()[0])) try: BOT_TOKEN = getConfig('BOT_TOKEN') parent_id = getConfig('GDRIVE_FOLDER_ID') diff --git a/bot/helper/telegram_helper/bot_commands.py b/bot/helper/telegram_helper/bot_commands.py index 7db7d57..be1833c 100644 --- a/bot/helper/telegram_helper/bot_commands.py +++ b/bot/helper/telegram_helper/bot_commands.py @@ -8,6 +8,7 @@ class _BotCommands: self.ListCommand = 'list' self.StatusCommand = 'status' self.AuthorizeCommand = 'authorize' + self.UnAuthorizeCommand = 'unauthorize' self.PingCommand = 'ping' self.DiskCommand = 'disk' self.HelpCommand = 'help' diff --git a/bot/modules/authorize.py b/bot/modules/authorize.py index 512894b..5c94053 100644 --- a/bot/modules/authorize.py +++ b/bot/modules/authorize.py @@ -7,31 +7,63 @@ from telegram.ext import Filters from telegram import Update from bot.helper.telegram_helper.bot_commands import BotCommands + @run_async def authorize(update: Update, context): reply_message = update.message.reply_to_message msg = '' with open('authorized_chats.txt', 'a') as file: if reply_message is None: + # Trying to authorize a chat chat_id = update.effective_chat.id if chat_id not in AUTHORIZED_CHATS: file.write(f'{chat_id}\n') - AUTHORIZED_CHATS.append(chat_id) + AUTHORIZED_CHATS.add(chat_id) msg = 'Chat authorized' else: msg = 'Already authorized chat' else: + # Trying to authorize someone in specific user_id = reply_message.from_user.id if user_id not in AUTHORIZED_CHATS: file.write(f'{user_id}\n') - AUTHORIZED_CHATS.append(user_id) + AUTHORIZED_CHATS.add(user_id) msg = 'Person Authorized to use the bot!' else: msg = 'Person already authorized' sendMessage(msg, context, update) +@run_async +def unauthorize(update, context): + reply_message = update.message.reply_to_message + if reply_message is None: + # Trying to unauthorize a chat + chat_id = update.effective_chat.id + if chat_id in AUTHORIZED_CHATS: + AUTHORIZED_CHATS.remove(chat_id) + msg = 'Chat unauthorized' + else: + msg = 'Already unauthorized chat' + else: + # Trying to authorize someone in specific + user_id = reply_message.from_user.id + if user_id in AUTHORIZED_CHATS: + AUTHORIZED_CHATS.remove(user_id) + msg = 'Person unauthorized to use the bot!' + else: + msg = 'Person already unauthorized!' + with open('authorized_chats.txt', 'a') as file: + file.truncate(0) + for i in AUTHORIZED_CHATS: + file.write(f'{i}\n') + sendMessage(msg, context, update) + + authorize_handler = CommandHandler(command=BotCommands.AuthorizeCommand, callback=authorize, filters=CustomFilters.owner_filter & Filters.group) +unauthorize_handler = CommandHandler(command=BotCommands.UnAuthorizeCommand, callback=unauthorize, + filters=CustomFilters.owner_filter & Filters.group) dispatcher.add_handler(authorize_handler) +dispatcher.add_handler(unauthorize_handler)