From f4290210b12b644d07a46f8244e1c7d9bfe7ec8d Mon Sep 17 00:00:00 2001 From: lzzy12 Date: Thu, 14 May 2020 19:02:33 +0530 Subject: [PATCH] Implemented mirror ownership Signed-off-by: lzzy12 --- bot/__init__.py | 2 +- bot/helper/telegram_helper/filters.py | 22 +++++++++++++++++++++- bot/modules/cancel_mirror.py | 16 ++++++++-------- bot/modules/mirror.py | 1 - 4 files changed, 30 insertions(+), 11 deletions(-) diff --git a/bot/__init__.py b/bot/__init__.py index 1f5cf40..3097979 100644 --- a/bot/__init__.py +++ b/bot/__init__.py @@ -54,7 +54,7 @@ status_reply_dict_lock = threading.Lock() # Value: telegram.Message status_reply_dict = {} # Key: update.message.message_id -# Value: An object of DownloadStatus +# Value: An object of Status download_dict = {} # Stores list of users and chats the bot is authorized to use in AUTHORIZED_CHATS = set() diff --git a/bot/helper/telegram_helper/filters.py b/bot/helper/telegram_helper/filters.py index a65445f..c850d6f 100644 --- a/bot/helper/telegram_helper/filters.py +++ b/bot/helper/telegram_helper/filters.py @@ -1,5 +1,6 @@ from telegram.ext import BaseFilter -from bot import AUTHORIZED_CHATS, OWNER_ID +from telegram import Message +from bot import AUTHORIZED_CHATS, OWNER_ID, download_dict, download_dict_lock class CustomFilters: @@ -21,3 +22,22 @@ class CustomFilters: return bool(message.chat.id in AUTHORIZED_CHATS) authorized_chat = _AuthorizedChat() + + class _MirrorOwner(BaseFilter): + def filter(self, message: Message): + user_id = message.from_user.id + if user_id == OWNER_ID: + return True + args = str(message.text).split(' ') + if len(args) > 1: + # Cancelling by gid + with download_dict_lock: + for message_id, status in download_dict.items(): + if status.gid == args[1] and status.message.from_user.id == user_id: + return True + else: + return False + # Cancelling by replying to original mirror message + reply_user = message.reply_to_message.from_user.id + return bool(reply_user == user_id) + mirror_owner_filter = _MirrorOwner() diff --git a/bot/modules/cancel_mirror.py b/bot/modules/cancel_mirror.py index c25c460..a4d2d74 100644 --- a/bot/modules/cancel_mirror.py +++ b/bot/modules/cancel_mirror.py @@ -11,14 +11,14 @@ from bot.helper.ext_utils.bot_utils import getDownloadByGid, MirrorStatus @run_async -def cancel_mirror(update,context): - args = update.message.text.split(" ",maxsplit=1) +def cancel_mirror(update, context): + args = update.message.text.split(" ", maxsplit=1) mirror_message = None if len(args) > 1: gid = args[1] dl = getDownloadByGid(gid) if not dl: - sendMessage(f"GID: {gid} not found.",context.bot,update) + sendMessage(f"GID: {gid} not found.", context.bot, update) return with download_dict_lock: keys = list(download_dict.keys()) @@ -33,11 +33,11 @@ def cancel_mirror(update,context): if BotCommands.MirrorCommand in mirror_message.text or \ BotCommands.TarMirrorCommand in mirror_message.text: msg = "Mirror already have been cancelled" - sendMessage(msg,context.bot,update) + sendMessage(msg, context.bot, update) return else: msg = "Please reply to the /mirror message which was used to start the download or /cancel gid to cancel it!" - sendMessage(msg,context.bot,update) + sendMessage(msg, context.bot, update) return if dl.status() == "Uploading": sendMessage("Upload in Progress, Don't Cancel it.", context.bot, update) @@ -56,16 +56,16 @@ def cancel_all(update, context): with download_dict_lock: count = 0 for dlDetails in list(download_dict.values()): - if dlDetails.status() == MirrorStatus.STATUS_DOWNLOADING\ + if dlDetails.status() == MirrorStatus.STATUS_DOWNLOADING \ or dlDetails.status() == MirrorStatus.STATUS_WAITING: dlDetails.download().cancel_download() count += 1 delete_all_messages() - sendMessage(f'Cancelled {count} downloads!', context.bot,update) + sendMessage(f'Cancelled {count} downloads!', context.bot, update) cancel_mirror_handler = CommandHandler(BotCommands.CancelMirror, cancel_mirror, - filters=CustomFilters.authorized_chat | CustomFilters.authorized_user) + filters=(CustomFilters.authorized_chat | CustomFilters.authorized_user) & CustomFilters.mirror_owner_filter) cancel_all_handler = CommandHandler(BotCommands.CancelAllCommand, cancel_all, filters=CustomFilters.owner_filter) dispatcher.add_handler(cancel_all_handler) diff --git a/bot/modules/mirror.py b/bot/modules/mirror.py index 0295a94..0c613c5 100644 --- a/bot/modules/mirror.py +++ b/bot/modules/mirror.py @@ -16,7 +16,6 @@ from bot.helper.mirror_utils.upload_utils import gdriveTools from bot.helper.telegram_helper.bot_commands import BotCommands from bot.helper.telegram_helper.filters import CustomFilters from bot.helper.telegram_helper.message_utils import * -from bot.helper.mirror_utils.download_utils.youtube_dl_download_helper import YoutubeDLHelper import pathlib import os