diff --git a/bot/__main__.py b/bot/__main__.py index 3bdd7c6..a94389f 100644 --- a/bot/__main__.py +++ b/bot/__main__.py @@ -6,7 +6,7 @@ import time from bot.helper.telegram_helper.message_utils import * import shutil from .helper.telegram_helper.filters import CustomFilters -from .modules import * +from .modules import authorize, list, cancel_mirror, mirror_status, mirror @run_async @@ -39,15 +39,22 @@ def ping(update, context): @run_async def bot_help(update, context): help_string = ''' - /help: To get this message +/help: To get this message - /mirror [download_url][magnet_link]: Start mirroring the link to google drive +/mirror [download_url][magnet_link]: Start mirroring the link to google drive - /cancel: Reply to the message by which the download was initiated and that download will be cancelled +/tarmirror [download_url][magnet_link]: start mirroring and upload the archived (.tar) version of the download - /status: Shows a status of all the downloads +/cancel: Reply to the message by which the download was initiated and that download will be cancelled - /list [search term]: Searches the search term in the Google drive, if found replies with the link''' +/status: Shows a status of all the downloads + +/list [search term]: Searches the search term in the Google drive, if found replies with the link + +/disk: Show a status of the disk usage of the machine the bot is hosted on + +/authorize: Authorize a chat or a user to use the bot (Can only be invoked by owner of the bot) +''' sendMessage(help_string, context, update) diff --git a/bot/helper/ext_utils/bot_utils.py b/bot/helper/ext_utils/bot_utils.py index fb7dc04..454e178 100644 --- a/bot/helper/ext_utils/bot_utils.py +++ b/bot/helper/ext_utils/bot_utils.py @@ -84,5 +84,5 @@ def get_readable_message(progress_list: list = download_dict.values()): f'{get_progress_bar_string(status)} {status.progress()} of {status.size()}\n' \ f'Speed: {status.speed()}\n' \ f'ETA: {status.eta()}\n' - LOGGER.info(msg) +# LOGGER.info(msg) return msg diff --git a/bot/helper/ext_utils/fs_utils.py b/bot/helper/ext_utils/fs_utils.py index 71f3a6f..46f9517 100644 --- a/bot/helper/ext_utils/fs_utils.py +++ b/bot/helper/ext_utils/fs_utils.py @@ -27,8 +27,9 @@ def exit_clean_up(signal, frame): def tar(orig_path: str): path = pathlib.PurePath(orig_path) base = path.name - root = path.parent.name - shutil.make_archive(base_name=orig_path, format='tar', root_dir=root, base_dir=base) + root = pathlib.Path(path.parent.as_posix()).absolute().as_posix() + LOGGER.info(f'Tar: orig_path: {orig_path}, base: {base}, root: {root}') + return shutil.make_archive(orig_path, 'tar', root, base) def get_mime_type(file_path): diff --git a/bot/helper/mirror_utils/download_status.py b/bot/helper/mirror_utils/download_status.py index 8c6e489..010a895 100644 --- a/bot/helper/mirror_utils/download_status.py +++ b/bot/helper/mirror_utils/download_status.py @@ -9,6 +9,7 @@ def get_download(gid): class DownloadStatus: def __init__(self, gid, message_id): + self.upload_name = None self.__gid = gid self.__download = get_download(gid) self.__uid = message_id @@ -55,6 +56,8 @@ class DownloadStatus: return self.__download.download_speed_string() def name(self): + if self.upload_name is not None: + return self.upload_name return self.__download.name def path(self): diff --git a/bot/modules/mirror.py b/bot/modules/mirror.py index dc621ec..1b03abf 100644 --- a/bot/modules/mirror.py +++ b/bot/modules/mirror.py @@ -1,17 +1,19 @@ from telegram.ext import CommandHandler, run_async from telegram.error import BadRequest from bot.helper.mirror_utils import download_tools, gdriveTools, listeners -from bot import LOGGER, dispatcher +from bot import LOGGER, dispatcher, DOWNLOAD_DIR from bot.helper.ext_utils import fs_utils from bot import download_dict, status_reply_dict, status_reply_dict_lock, download_dict_lock from bot.helper.telegram_helper.message_utils import * 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 class MirrorListener(listeners.MirrorListeners): - def __init__(self, context, update, reply_message): + def __init__(self, context, update, reply_message, isTar=False): super().__init__(context, update, reply_message) + self.isTar = isTar def onDownloadStarted(self, link): LOGGER.info("Adding link: " + link) @@ -28,8 +30,14 @@ class MirrorListener(listeners.MirrorListeners): def onDownloadComplete(self, progress_status_list, index: int): LOGGER.info(f"Download completed: {progress_status_list[index].name()}") + if self.isTar: + path = fs_utils.tar(f'{DOWNLOAD_DIR}{self.uid}/{progress_status_list[index].name()}') + else: + path = f'{DOWNLOAD_DIR}{self.uid}/{progress_status_list[index].name()}' + name = pathlib.PurePath(path).name + download_dict[self.uid].upload_name = name gdrive = gdriveTools.GoogleDriveHelper(self) - gdrive.upload(progress_status_list[index].name()) + gdrive.upload(name) def onDownloadError(self, error, progress_status_list: list, index: int): LOGGER.error(error) @@ -78,10 +86,10 @@ class MirrorListener(listeners.MirrorListeners): raise KillThreadException('Message deleted. Do not call this method from the thread') -@run_async -def mirror(update, context): - message = update.message.text - link = message.replace('/mirror', '')[1:] +def _mirror(update, context, isTar=False): + message_args = update.message.text.split(' ') + link = message_args[1] + 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 @@ -93,10 +101,23 @@ def mirror(update, context): if index in status_reply_dict.keys(): deleteMessage(context, status_reply_dict[index]) status_reply_dict[index] = reply_msg - listener = MirrorListener(context, update, reply_msg) + listener = MirrorListener(context, update, reply_msg, isTar) aria = download_tools.DownloadHelper(listener) aria.add_download(link) +@run_async +def mirror(update, context): + _mirror(update, context) + + +@run_async +def tar_mirror(update, context): + _mirror(update, context, True) + + mirror_handler = CommandHandler('mirror', mirror, filters=CustomFilters.authorized_chat | CustomFilters.authorized_user) +tar_mirror_handler = CommandHandler('tarmirror', tar_mirror, + filters=CustomFilters.authorized_chat | CustomFilters.authorized_user) dispatcher.add_handler(mirror_handler) +dispatcher.add_handler(tar_mirror_handler)