diff --git a/.gitignore b/.gitignore index 7b7e0bf..04bfd67 100644 --- a/.gitignore +++ b/.gitignore @@ -11,3 +11,4 @@ data* .idea *.json *.pickle +authorized_chats.txt \ No newline at end of file diff --git a/bot/__init__.py b/bot/__init__.py index f22c63d..a529f0f 100644 --- a/bot/__init__.py +++ b/bot/__init__.py @@ -39,6 +39,13 @@ download_dict_lock = threading.Lock() status_reply_dict_lock = threading.Lock() status_reply_dict = {} download_dict = {} +# Stores list of users and chats the bot is authorized to use in +AUTHORIZED_CHATS = [] +with open('authorized_chats.txt', 'r+') as f: + lines = f.readlines() + for line in lines: + LOGGER.info(line.split()) + AUTHORIZED_CHATS.append(line.split()[0]) try: BOT_TOKEN = getConfig('BOT_TOKEN') parent_id = getConfig('GDRIVE_FOLDER_ID') @@ -46,7 +53,7 @@ try: if DOWNLOAD_DIR[-1] != '/' or DOWNLOAD_DIR[-1] != '\\': DOWNLOAD_DIR = DOWNLOAD_DIR + '/' DOWNLOAD_STATUS_UPDATE_INTERVAL = int(getConfig('DOWNLOAD_STATUS_UPDATE_INTERVAL')) - + OWNER_ID = int(getConfig('OWNER_ID')) except KeyError as e: LOGGER.error("One or more env variables missing! Exiting now") exit(1) diff --git a/bot/__main__.py b/bot/__main__.py index 1e069de..92b03fe 100644 --- a/bot/__main__.py +++ b/bot/__main__.py @@ -1,12 +1,12 @@ from telegram.ext import CommandHandler, run_async from bot import dispatcher, LOGGER, updater, aria2 -import bot.mirror, bot.list, bot.mirror_status, bot.cancel_mirror +import bot.mirror, bot.list, bot.mirror_status, bot.cancel_mirror, bot.authorize from bot.helper import fs_utils import signal import time from bot.helper.message_utils import * import shutil - +from .helper.telegram_helper.filters import CustomFilters @run_async def disk_usage(update, context): @@ -52,14 +52,19 @@ def bot_help(update, context): def main(): - start_handler = CommandHandler('start', start) - ping_handler = CommandHandler('ping', ping) - help_handler = CommandHandler('help', bot_help) - disk_usage_handler = CommandHandler('disk', disk_usage) + start_handler = CommandHandler('start', start, + filters=CustomFilters.authorized_chat | CustomFilters.authorized_user) + ping_handler = CommandHandler('ping', ping, + filters=CustomFilters.authorized_chat | CustomFilters.authorized_user) + help_handler = CommandHandler('help', + bot_help, filters=CustomFilters.authorized_chat | CustomFilters.authorized_user) + disk_handler = CommandHandler('disk', + disk_usage, filters=CustomFilters.authorized_chat | CustomFilters.authorized_user) + dispatcher.add_handler(start_handler) dispatcher.add_handler(ping_handler) dispatcher.add_handler(help_handler) - dispatcher.add_handler(disk_usage_handler) + dispatcher.add_handler(disk_handler) updater.start_polling() LOGGER.info("Bot Started!") signal.signal(signal.SIGINT, fs_utils.exit_clean_up) diff --git a/bot/authorize.py b/bot/authorize.py new file mode 100644 index 0000000..f3d85b9 --- /dev/null +++ b/bot/authorize.py @@ -0,0 +1,32 @@ +from .helper.message_utils import sendMessage +from telegram.ext import run_async +from bot import AUTHORIZED_CHATS, dispatcher +from telegram.ext import CommandHandler +from .helper.telegram_helper.filters import CustomFilters +from telegram.ext import Filters +from telegram import Update + + +@run_async +def authorize(update: Update, context): + reply_message = update.message.reply_to_message + with open('authorized_chats.txt', 'a') as file: + if reply_message is None: + chat_id = update.effective_chat.id + if chat_id not in AUTHORIZED_CHATS: + file.write(f'{chat_id}\n') + AUTHORIZED_CHATS.append(chat_id) + sendMessage('Chat authorized', context, update) + else: + sendMessage('Already authorized chat', context, update) + else: + 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) + + +authorize_handler = CommandHandler(command='authorize', callback=authorize, + filters=CustomFilters.owner_filter & Filters.group) +dispatcher.add_handler(authorize_handler) + diff --git a/bot/cancel_mirror.py b/bot/cancel_mirror.py index fdc38b0..f4575a8 100644 --- a/bot/cancel_mirror.py +++ b/bot/cancel_mirror.py @@ -1,7 +1,7 @@ from telegram.ext import CommandHandler, run_async from bot.helper.message_utils import * from bot import download_dict, aria2, dispatcher, download_dict_lock - +from .helper.telegram_helper.filters import CustomFilters @run_async def cancel_mirror(update: Update, context): @@ -21,5 +21,6 @@ def cancel_mirror(update: Update, context): sendMessage("Download canceled", context, update) -cancel_mirror_handler = CommandHandler('cancel', cancel_mirror) -dispatcher.add_handler(cancel_mirror_handler) \ No newline at end of file +cancel_mirror_handler = CommandHandler('cancel', cancel_mirror, + filters=CustomFilters.authorized_chat | CustomFilters.authorized_user) +dispatcher.add_handler(cancel_mirror_handler) diff --git a/bot/config_sample.ini b/bot/config_sample.ini index 1e08754..586c84c 100644 --- a/bot/config_sample.ini +++ b/bot/config_sample.ini @@ -6,4 +6,5 @@ _____REMOVE_THIS_LINE_____=True BOT_TOKEN = GDRIVE_FOLDER_ID = DOWNLOAD_DIR = ./download -DOWNLOAD_STATUS_UPDATE_INTERVAL = 5 \ No newline at end of file +DOWNLOAD_STATUS_UPDATE_INTERVAL = 5 +OWNER_ID = \ No newline at end of file diff --git a/bot/helper/telegram_helper/__init__.py b/bot/helper/telegram_helper/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/bot/helper/telegram_helper/filters.py b/bot/helper/telegram_helper/filters.py new file mode 100644 index 0000000..0083f4a --- /dev/null +++ b/bot/helper/telegram_helper/filters.py @@ -0,0 +1,24 @@ +from telegram.ext import BaseFilter +from telegram import Update +from bot import AUTHORIZED_CHATS, OWNER_ID + + +class CustomFilters: + class _OwnerFilter(BaseFilter): + def filter(self, message): + return bool(message.from_user.id == OWNER_ID) + + owner_filter = _OwnerFilter() + + class _AuthorizedUserFilter(BaseFilter): + def filter(self, message): + id = message.from_user.id + return bool(id in AUTHORIZED_CHATS or id == OWNER_ID) + + authorized_user = _AuthorizedUserFilter() + + class _AuthorizedChat(BaseFilter): + def filter(self, message): + return bool(message.chat.id in AUTHORIZED_CHATS) + + authorized_chat = _AuthorizedChat() diff --git a/bot/mirror.py b/bot/mirror.py index 893795d..4c83a32 100644 --- a/bot/mirror.py +++ b/bot/mirror.py @@ -7,7 +7,7 @@ from bot import download_dict, status_reply_dict, status_reply_dict_lock, downlo from bot.helper.message_utils import * from bot.helper.bot_utils import get_readable_message, MirrorStatus from bot.helper.exceptions import KillThreadException - +from .helper.telegram_helper.filters import CustomFilters class MirrorListener(listeners.MirrorListeners): def __init__(self, context, update, reply_message): @@ -98,5 +98,5 @@ def mirror(update, context): aria.add_download(link) -mirror_handler = CommandHandler('mirror', mirror) +mirror_handler = CommandHandler('mirror', mirror, filters=CustomFilters.authorized_chat | CustomFilters.authorized_user) dispatcher.add_handler(mirror_handler) diff --git a/bot/mirror_status.py b/bot/mirror_status.py index f3b738c..903d55b 100644 --- a/bot/mirror_status.py +++ b/bot/mirror_status.py @@ -4,7 +4,7 @@ from bot.helper.message_utils import * from time import sleep from bot.helper.bot_utils import get_readable_message from telegram.error import BadRequest - +from .helper.telegram_helper.filters import CustomFilters @run_async def mirror_status(update: Update, context): @@ -42,5 +42,6 @@ def mirror_status(update: Update, context): sleep(DOWNLOAD_STATUS_UPDATE_INTERVAL) -mirror_status_handler = CommandHandler('status', mirror_status) +mirror_status_handler = CommandHandler('status', mirror_status, + filters=CustomFilters.authorized_chat | CustomFilters.authorized_user) dispatcher.add_handler(mirror_status_handler)