From 9d83cd31c705c1076438f2216c40d38a492deb12 Mon Sep 17 00:00:00 2001 From: lzzy12 Date: Fri, 11 Oct 2019 00:07:46 +0530 Subject: [PATCH] Random fixes all around the source Signed-off-by: lzzy12 --- bot/__init__.py | 24 ++++++++++++----------- bot/helper/ext_utils/exceptions.py | 8 ++++++++ bot/helper/mirror_utils/download_tools.py | 15 +++++++------- bot/modules/cancel_mirror.py | 10 +++++++--- bot/modules/mirror.py | 21 +++++++++++++++----- bot/modules/mirror_status.py | 1 + 6 files changed, 52 insertions(+), 27 deletions(-) diff --git a/bot/__init__.py b/bot/__init__.py index fd76c5c..eeddf25 100644 --- a/bot/__init__.py +++ b/bot/__init__.py @@ -3,6 +3,7 @@ import configparser import aria2p import threading from telegram.ext import Updater +import os logging.basicConfig(format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', level=logging.INFO) @@ -25,12 +26,12 @@ except KeyError: pass aria2 = aria2p.API( - aria2p.Client( - host="http://localhost", - port=6800, - secret="", - ) - ) + aria2p.Client( + host="http://localhost", + port=6800, + secret="", + ) +) DOWNLOAD_DIR = None BOT_TOKEN = None @@ -41,11 +42,12 @@ 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(int(line.split()[0])) +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])) try: BOT_TOKEN = getConfig('BOT_TOKEN') parent_id = getConfig('GDRIVE_FOLDER_ID') diff --git a/bot/helper/ext_utils/exceptions.py b/bot/helper/ext_utils/exceptions.py index 77f0ccf..e39e15a 100644 --- a/bot/helper/ext_utils/exceptions.py +++ b/bot/helper/ext_utils/exceptions.py @@ -4,6 +4,14 @@ class DriveAuthError(Exception): class KillThreadException(Exception): """ Custom Exception class for killing thread as soon as they aren't needed""" + + def __init__(self, message, error=None): + super().__init__(message) + self.error = error + + +class DownloadCancelled(Exception): + def __init__(self, message, error=None): super().__init__(message) self.error = error diff --git a/bot/helper/mirror_utils/download_tools.py b/bot/helper/mirror_utils/download_tools.py index 4636c2e..2351fcc 100644 --- a/bot/helper/mirror_utils/download_tools.py +++ b/bot/helper/mirror_utils/download_tools.py @@ -24,13 +24,12 @@ class DownloadHelper: return False def add_download(self, link: str): - download = None if self.is_url(link): if link.endswith('.torrent'): self.__is_torrent = True - download = aria2.add_uris([link], {'dir': DOWNLOAD_DIR + str(self.__listener.message.message_id)}) + download = aria2.add_uris([link], {'dir': DOWNLOAD_DIR + str(self.__listener.uid)}) elif self.is_magnet(link): - download = aria2.add_magnet(link, {'dir': DOWNLOAD_DIR + str(self.__listener.message.message_id)}) + download = aria2.add_magnet(link, {'dir': DOWNLOAD_DIR + str(self.__listener.uid)}) self.__is_torrent = True else: _list = get_download_status_list() @@ -38,12 +37,12 @@ class DownloadHelper: return with download_dict_lock: download_dict[self.__listener.message.message_id] = DownloadStatus(download.gid, - self.__listener.message.message_id) + self.__listener.uid) self.__listener.onDownloadStarted(link) self.__update_download_status() def __get_download(self): - return get_download(self.__listener.message.message_id) + return get_download(self.__listener.uid) def __get_followed_download_gid(self): download = self.__get_download() @@ -59,8 +58,8 @@ class DownloadHelper: if self.__is_torrent: # Waiting for the actual gid new_gid = None - download = self.__get_download() while new_gid is None: + download = self.__get_download() if download.has_failed: self.__listener.onDownloadError(download.error_message, status_list, index) return @@ -69,12 +68,12 @@ class DownloadHelper: return sleep(DOWNLOAD_STATUS_UPDATE_INTERVAL) if should_update: - # Check every few seconds - new_gid = self.__get_followed_download_gid() try: self.__listener.onDownloadProgress(get_download_status_list(), index) except KillThreadException: should_update = False + # Check every few seconds + new_gid = self.__get_followed_download_gid() with download_dict_lock: download_dict[self.__listener.message.message_id] = DownloadStatus(new_gid, self.__listener.message.message_id) diff --git a/bot/modules/cancel_mirror.py b/bot/modules/cancel_mirror.py index 365a9ae..20721cf 100644 --- a/bot/modules/cancel_mirror.py +++ b/bot/modules/cancel_mirror.py @@ -1,7 +1,9 @@ from telegram.ext import CommandHandler, run_async from bot.helper.telegram_helper.message_utils import * -from bot import download_dict, aria2, dispatcher, download_dict_lock +from bot import download_dict, aria2, dispatcher, download_dict_lock, DOWNLOAD_DIR from bot.helper.telegram_helper.filters import CustomFilters +from bot.helper.ext_utils.fs_utils import clean_download + @run_async def cancel_mirror(update: Update, context): @@ -10,13 +12,15 @@ def cancel_mirror(update: Update, context): keys = download_dict.keys() download = download_dict[mirror_message.message_id].download() if mirror_message is None or mirror_message.message_id not in keys: - if '/mirror' in mirror_message.text: + if '/mirror' in mirror_message.text or '/tarmirror' in mirror_message.text: msg = 'Message has already been cancelled' else: msg = 'Please reply to the /mirror message which was used to start the download to cancel it!' sendMessage(msg, context, update) return - + if len(download.followed_by_ids) != 0: + downloads = aria2.get_downloads(download.followed_by_ids) + aria2.pause(downloads) aria2.pause([download]) sendMessage("Download canceled", context, update) diff --git a/bot/modules/mirror.py b/bot/modules/mirror.py index dd9a2c3..0d36e1b 100644 --- a/bot/modules/mirror.py +++ b/bot/modules/mirror.py @@ -9,6 +9,8 @@ from bot.helper.ext_utils.bot_utils import get_readable_message, MirrorStatus from bot.helper.ext_utils.exceptions import KillThreadException from bot.helper.telegram_helper.filters import CustomFilters import pathlib +import threading + class MirrorListener(listeners.MirrorListeners): def __init__(self, context, update, reply_message, isTar=False): @@ -19,9 +21,10 @@ class MirrorListener(listeners.MirrorListeners): LOGGER.info("Adding link: " + link) def onDownloadProgress(self, progress_status_list: list, index: int): - if progress_status_list[index].status() == MirrorStatus.STATUS_CANCELLED: - raise KillThreadException('Mirror cancelled by user') msg = get_readable_message(progress_status_list) + if progress_status_list[index].status() == MirrorStatus.STATUS_CANCELLED: + editMessage(msg, self.context, self.reply_message) + raise KillThreadException('Mirror cancelled by user') # LOGGER.info("Editing message") try: editMessage(msg, self.context, self.reply_message) @@ -50,9 +53,9 @@ class MirrorListener(listeners.MirrorListeners): deleteMessage(self.context, status_reply_dict[self.update.effective_chat.id]) del status_reply_dict[self.update.effective_chat.id] if index is not None: - fs_utils.clean_download(progress_status_list[index].path()) with download_dict_lock: del download_dict[self.message.message_id] + fs_utils.clean_download(progress_status_list[index].path()) msg = f"@{self.message.from_user.username} your download has been stopped due to: {error}" sendMessage(msg, self.context, self.update) @@ -92,13 +95,20 @@ class MirrorListener(listeners.MirrorListeners): def _mirror(update, context, isTar=False): message_args = update.message.text.split(' ') - link = message_args[1] + try: + link = message_args[1] + except KeyError: + link = '' LOGGER.info(link) link = link.strip() + if len(link) == 0 and update.message.reply_to_message is not None: document = update.message.reply_to_message.document if document is not None and document.mime_type == "application/x-bittorrent": link = document.get_file().file_path + else: + sendMessage('No download source provided', context, update) + return reply_msg = sendMessage('Starting Download', context, update) index = update.effective_chat.id with status_reply_dict_lock: @@ -107,7 +117,8 @@ def _mirror(update, context, isTar=False): status_reply_dict[index] = reply_msg listener = MirrorListener(context, update, reply_msg, isTar) aria = download_tools.DownloadHelper(listener) - aria.add_download(link) + t = threading.Thread(target=aria.add_download, args=(link,)) + t.start() @run_async diff --git a/bot/modules/mirror_status.py b/bot/modules/mirror_status.py index e812c1e..60776ff 100644 --- a/bot/modules/mirror_status.py +++ b/bot/modules/mirror_status.py @@ -6,6 +6,7 @@ from bot.helper.ext_utils.bot_utils import get_readable_message from telegram.error import BadRequest from bot.helper.telegram_helper.filters import CustomFilters + @run_async def mirror_status(update: Update, context): message = get_readable_message()