diff --git a/bot/__init__.py b/bot/__init__.py index e7ea0c3..38fc273 100644 --- a/bot/__init__.py +++ b/bot/__init__.py @@ -7,7 +7,7 @@ logging.basicConfig(format='%(asctime)s - %(name)s - %(levelname)s - %(message)s level=logging.INFO) config = configparser.ConfigParser() -config.read('bot/conf.ini') +config.read('bot/config.ini') def getConfig(name: str): diff --git a/bot/helper/download_status.py b/bot/helper/download_status.py index 5278632..171f772 100644 --- a/bot/helper/download_status.py +++ b/bot/helper/download_status.py @@ -1,4 +1,4 @@ -from bot import aria2 +from bot import aria2, download_list, DOWNLOAD_DIR def get_download(gid): @@ -11,9 +11,10 @@ class DownloadStatus: STATUS_WAITING = "Queued" STATUS_FAILED = "Failed. Cleaning download" - def __init__(self, gid): + def __init__(self, gid, update_id): self.__gid = gid self.__download = get_download(gid) + self.__uid = update_id def __update(self): self.__download = get_download(self.__gid) @@ -29,6 +30,9 @@ class DownloadStatus: def name(self): return self.__download.name + def path(self): + return "{}{}/{}".format(DOWNLOAD_DIR, self.__uid, self.name()) + def size(self): return self.__download.total_length_string() @@ -45,6 +49,8 @@ class DownloadStatus: # If download exists and is complete the it must be uploading # otherwise the gid would have been removed from the download_list status = DownloadStatus.STATUS_UPLOADING + elif self.__download.has_failed: + status = DownloadStatus.STATUS_FAILED elif self.__download.is_active: status = DownloadStatus.STATUS_DOWNLOADING return status diff --git a/bot/helper/download_tools.py b/bot/helper/download_tools.py index 78b669c..a779175 100644 --- a/bot/helper/download_tools.py +++ b/bot/helper/download_tools.py @@ -1,13 +1,12 @@ from time import sleep -from bot import LOGGER, DOWNLOAD_DIR, DOWNLOAD_STATUS_UPDATE_INTERVAL, download_list, aria2 -from bot.helper.listeners import * +from bot import DOWNLOAD_DIR, DOWNLOAD_STATUS_UPDATE_INTERVAL, aria2 from .download_status import DownloadStatus from .bot_utils import * class DownloadHelper: - def __init__(self, listener: MirrorListeners): + def __init__(self, listener=None): self.__listener = listener self.__is_torrent = False @@ -33,7 +32,7 @@ class DownloadHelper: else: self.__listener.onDownloadError("No download URL or URL malformed") return - download_list[self.__listener.update.update_id] = DownloadStatus(download.gid) + download_list[self.__listener.update.update_id] = DownloadStatus(download.gid, self.__listener.update.update_id) self.__listener.onDownloadStarted(link) self.__update_download_status() @@ -57,17 +56,18 @@ class DownloadHelper: sleep(DOWNLOAD_STATUS_UPDATE_INTERVAL) new_gid = self.__get_followed_download_gid() self.__listener.onDownloadProgress(get_download_status_list(), index) - download_list[self.__listener.update.update_id] = DownloadStatus(new_gid) + download_list[self.__listener.update.update_id] = DownloadStatus(new_gid, self.__listener.update.update_id) # Start tracking the actual download previous = None download = self.__get_download() while not download.is_complete: status_list = get_download_status_list() + index = get_download_index(status_list, self.__get_download().gid) if self.__get_download().has_failed: - self.__listener.onDownloadError(self.__get_download().error_message) + self.__listener.onDownloadError(self.__get_download().error_message, status_list[index]) break - # TODO: Find a better way + # TODO: Find a better way to differentiate between 2 list of objects progress_str_list = get_download_str() if progress_str_list != previous: self.__listener.onDownloadProgress(status_list, index) diff --git a/bot/helper/fs_utils.py b/bot/helper/fs_utils.py index a16f01d..bd840a7 100644 --- a/bot/helper/fs_utils.py +++ b/bot/helper/fs_utils.py @@ -5,6 +5,7 @@ import os import pathlib import mimetypes + def clean_download(path: str): if os.path.exists(path): shutil.rmtree(path) @@ -13,8 +14,8 @@ def clean_download(path: str): def exit_clean_up(signal, frame): try: LOGGER.info("Please wait, while we clean up the downloads and stop running downloads") - aria2.autopurge() - aria2.remove_all(True) + for download in aria2.get_downloads(): + download.remove() shutil.rmtree(DOWNLOAD_DIR) sys.exit(0) except KeyboardInterrupt: diff --git a/bot/helper/gdriveTools.py b/bot/helper/gdriveTools.py index 7c1d7b0..59bc607 100644 --- a/bot/helper/gdriveTools.py +++ b/bot/helper/gdriveTools.py @@ -4,15 +4,14 @@ from google.auth.transport.requests import Request from googleapiclient.http import MediaFileUpload import pickle import os -from bot import LOGGER, parent_id, DOWNLOAD_DIR, download_list -from .listeners import MirrorListeners -from .fs_utils import clean_download, get_mime_type +from bot import LOGGER, parent_id, DOWNLOAD_DIR +from .fs_utils import get_mime_type from .bot_utils import * class GoogleDriveHelper: - def __init__(self, listener: MirrorListeners): + def __init__(self, listener=None): self.__G_DRIVE_TOKEN_FILE = "token.pickle" # Check https://developers.google.com/drive/scopes for all available scopes self.__OAUTH_SCOPE = "https://www.googleapis.com/auth/drive.file" @@ -63,28 +62,24 @@ class GoogleDriveHelper: if os.path.isfile(file_path): mime_type = get_mime_type(file_path) try: - g_drive_link = self.upload_file(file_path, file_name, mime_type,parent_id) + g_drive_link = self.upload_file(file_path, file_name, mime_type, parent_id) LOGGER.info("Uploaded To G-Drive: " + file_path) link = g_drive_link except Exception as e: LOGGER.error(str(e)) - pass + self.__listener.onUploadError(str(e), _list, index) else: try: - dir_id = self.create_directory(os.path.basename(os.path.abspath(file_name)),parent_id) + dir_id = self.create_directory(os.path.basename(os.path.abspath(file_name)), parent_id) self.upload_dir(file_path,dir_id) LOGGER.info("Uploaded To G-Drive: " + file_name) link = "https://drive.google.com/folderview?id={}".format(dir_id) except Exception as e: LOGGER.error(str(e)) - self.__listener.onUploadError(str(e)) - clean_download(file_dir) - raise Exception('Error: {}'.format(str(e))) - del download_list[self.__listener.update.update_id] + self.__listener.onUploadError(str(e), _list, index) LOGGER.info(download_list) self.__listener.onUploadComplete(link, _list, index) LOGGER.info("Deleting downloaded file/folder..") - clean_download(file_dir) return link def create_directory(self, directory_name,parent_id): diff --git a/bot/helper/listeners.py b/bot/helper/listeners.py index e6a7c39..3d6ac87 100644 --- a/bot/helper/listeners.py +++ b/bot/helper/listeners.py @@ -13,7 +13,7 @@ class MirrorListeners: def onDownloadComplete(self, progress_status_list: list, index: int): raise NotImplementedError - def onDownloadError(self, error: str): + def onDownloadError(self, error: str, progress_status_list: list, index: int): raise NotImplementedError def onUploadStarted(self, progress_status_list: list, index: int): @@ -22,5 +22,5 @@ class MirrorListeners: def onUploadComplete(self, link: str, progress_status_list: list, index: int): raise NotImplementedError - def onUploadError(self, error: str): + def onUploadError(self, error: str, progress_status: list, index: int): raise NotImplementedError \ No newline at end of file diff --git a/bot/helper/message_utils.py b/bot/helper/message_utils.py new file mode 100644 index 0000000..605463d --- /dev/null +++ b/bot/helper/message_utils.py @@ -0,0 +1,24 @@ +from telegram.ext import CommandHandler, run_async +from telegram.error import BadRequest +from telegram.message import Message +from telegram.update import Update + + +def sendMessage(text: str, context, update: Update): + return context.bot.send_message(update.message.chat_id, + reply_to_message_id=update.message.message_id, + text=text, parse_mode='HTMl') + + +def editMessage(text: str, context, message: Message): + try: + context.bot.edit_message_text(text=text, message_id=message.message_id, + chat_id=message.chat.id, + parse_mode='HTMl') + except BadRequest: + pass + + +def deleteMessage(context, message: Message): + context.bot.delete_message(chat_id=message.chat.id, + message_id=message.message_id) diff --git a/bot/mirror.py b/bot/mirror.py index ca5904b..4d2d673 100644 --- a/bot/mirror.py +++ b/bot/mirror.py @@ -1,8 +1,10 @@ -from telegram.ext import CommandHandler, run_async -from telegram.error import BadRequest + from bot.helper import download_tools, gdriveTools, listeners from bot import LOGGER, dispatcher - +from bot.helper import fs_utils +from bot.helper.download_status import DownloadStatus +from bot import download_list +from bot.helper.message_utils import * LOGGER.info('mirror.py') @@ -29,64 +31,51 @@ class MirrorListener(listeners.MirrorListeners): def onDownloadProgress(self, progress_status_list: list, index: int): msg = get_readable_message(progress_status_list) - if self.reply_message.text != msg: - LOGGER.info("Editing message") - self.context.bot.edit_message_text(text=msg, message_id=self.reply_message.message_id, - chat_id=self.reply_message.chat.id, - parse_mode='HTMl') + LOGGER.info("Editing message") + editMessage(msg, self.context, self.reply_message) def onDownloadComplete(self, progress_status_list, index: int): msg = get_readable_message(progress_status_list) LOGGER.info("Download completed") - try: - LOGGER.info("Editing message") - self.context.bot.edit_message_text(text=msg, message_id=self.reply_message.message_id, - chat_id=self.reply_message.chat.id, - parse_mode='HTMl') - except BadRequest: - pass + editMessage(msg, self.context, self.reply_message) gdrive = gdriveTools.GoogleDriveHelper(self) gdrive.upload(progress_status_list[index].name()) - def onDownloadError(self, error): + def onDownloadError(self, error, progress_status_list: list, index: int): LOGGER.error(error) - self.context.bot.edit_message_text(text=error, message_id=self.reply_message.message_id, - chat_id=self.reply_message.chat.id, - parse_mode='HTMl') + editMessage(error, self.context, self.reply_message) + fs_utils.clean_download(progress_status_list[index].path()) def onUploadStarted(self, progress_status_list: list, index: int): - try: - msg = get_readable_message(progress_status_list) - self.context.bot.edit_message_text(text=msg, message_id=self.reply_message.message_id, - chat_id=self.reply_message.chat.id, - parse_mode='HTMl') - except BadRequest: - pass + msg = get_readable_message(progress_status_list) + editMessage(msg, self.context, self.reply_message) def onUploadComplete(self, link: str, progress_status_list: list, index: int): msg = '{}'.format(link, progress_status_list[index].name()) - self.context.bot.delete_message(chat_id=self.reply_message.chat.id, message_id=self.reply_message.message_id) - self.context.bot.send_message(chat_id=self.update.message.chat_id, - reply_to_message_id=self.update.message.message_id, - text=msg, parse_mode='HTMl') + deleteMessage(self.context, self.reply_message) + sendMessage(msg, self.context, self.update) + fs_utils.clean_download(progress_status_list[index].path()) - def onUploadError(self, error: str): + def onUploadError(self, error: str, progress_status: list, index: int): LOGGER.error(error) - self.context.bot.edit_message_text(text=error, message_id=self.reply_message.message_id, - chat_id=self.reply_message.chat.id, - parse_mode='HTMl') + editMessage(error, self.context, self.reply_message) + fs_utils.clean_download(progress_status[index].path()) @run_async def mirror(update, context): message = update.message.text link = message.replace('/mirror', '')[1:] - reply_msg = context.bot.send_message(chat_id=update.message.chat_id, reply_to_message_id=update.message.message_id, - text="Starting Download") + reply_msg = sendMessage('Starting Download', context, update) listener = MirrorListener(context, update, reply_msg) aria = download_tools.DownloadHelper(listener) aria.add_download(link) +@run_async +def cancel_mirror(update, context): + pass + + mirror_handler = CommandHandler('mirror', mirror) dispatcher.add_handler(mirror_handler)