From 578eeb4004460e71f332f264f2b50cb43fd1e134 Mon Sep 17 00:00:00 2001 From: jaskaranSM Date: Thu, 28 May 2020 15:17:19 +0530 Subject: [PATCH] Implement Mega.nz mirror support. - Co-authored by - lzzy12 (github.com/lzzy12) --- bot/helper/ext_utils/bot_utils.py | 3 + .../download_utils/mega_downloader.py | 132 ++++++++++++++++++ .../status_utils/mega_download_status.py | 62 ++++++++ bot/modules/mirror.py | 7 +- 4 files changed, 203 insertions(+), 1 deletion(-) create mode 100644 bot/helper/mirror_utils/download_utils/mega_downloader.py create mode 100644 bot/helper/mirror_utils/status_utils/mega_download_status.py diff --git a/bot/helper/ext_utils/bot_utils.py b/bot/helper/ext_utils/bot_utils.py index a311ead..a015fb7 100644 --- a/bot/helper/ext_utils/bot_utils.py +++ b/bot/helper/ext_utils/bot_utils.py @@ -137,6 +137,9 @@ def is_magnet(url: str): return True return False +def is_mega_link(url: str): + return "mega.nz" in url + def new_thread(fn): """To use as decorator to make a function call threaded. Needs import diff --git a/bot/helper/mirror_utils/download_utils/mega_downloader.py b/bot/helper/mirror_utils/download_utils/mega_downloader.py new file mode 100644 index 0000000..4714b11 --- /dev/null +++ b/bot/helper/mirror_utils/download_utils/mega_downloader.py @@ -0,0 +1,132 @@ +from bot import LOGGER, MEGA_API_KEY, download_dict_lock, download_dict +import threading +from mega import (MegaApi, MegaListener, MegaRequest, MegaTransfer, MegaError) +from bot.helper.telegram_helper.message_utils import update_all_messages +import os +from bot.helper.mirror_utils.status_utils.mega_download_status import MegaDownloadStatus + +class MegaAppListener(MegaListener): + + def __init__(self, continue_event: threading.Event, listener): + self.continue_event = continue_event + self.node = None + self.listener = listener + self.uid = listener.uid + self.__bytes_transferred = 0 + self.is_cancelled = False + self.__speed = 0 + self.__name = '' + self.__size = 0 + self.error = None + self.gid = "" + super(MegaAppListener, self).__init__() + + @property + def speed(self): + """Returns speed of the download in bytes/second""" + return self.__speed + + @property + def name(self): + """Returns name of the download""" + return self.__name + + def setValues(self,name,size,gid): + self.__name = name + self.__size = size + self.gid = gid + + @property + def size(self): + """Size of download in bytes""" + return self.__size + + @property + def downloaded_bytes(self): + return self.__bytes_transferred + + def onRequestStart(self, api, request): + LOGGER.info('Request start ({})'.format(request)) + + def onRequestFinish(self, api, request, error): + LOGGER.info('Mega Request finished ({}); Result: {}' + .format(request, error)) + + request_type = request.getType() + if request_type == MegaRequest.TYPE_GET_PUBLIC_NODE: + self.node = request.getPublicMegaNode() + if request_type == MegaRequest.TYPE_LOGIN: + LOGGER.info("Fetching Nodes.") + api.fetchNodes() + if request_type == MegaRequest.TYPE_FETCH_NODES: + LOGGER.info("Fetching Root Node.") + self.node = api.getRootNode() + if request_type != MegaRequest.TYPE_LOGIN: + self.continue_event.set() + + def onRequestTemporaryError(self, api, request, error: MegaError): + self.listener.onDownloadError(error.toString()) + self.error = error.toString() + self.continue_event.set() + + def onTransferStart(self, api: MegaApi, transfer: MegaTransfer): + LOGGER.info(f"Transfer Started: {transfer.getFileName()}") + + def onTransferUpdate(self, api: MegaApi, transfer: MegaTransfer): + if self.is_cancelled: + api.cancelTransfer(transfer,None) + self.__speed = transfer.getSpeed() + self.__bytes_transferred = transfer.getTransferredBytes() + + def onTransferFinish(self, api: MegaApi, transfer : MegaTransfer, error): + try: + LOGGER.info(f'Transfer finished ({transfer}); Result: {transfer.getFileName()}') + if str(error) != "No error" and self.is_cancelled: + self.is_cancelled = False + return self.listener.onDownloadError(error.toString()) + if transfer.isFolderTransfer() and transfer.isFinished() and not self.is_cancelled or transfer.getFileName() == self.name and not self.is_cancelled: + self.listener.onDownloadComplete() + except Exception as e: + LOGGER.error(e) + + def onTransferTemporaryError(self, api, transfer, error): + LOGGER.info(f'Mega download error in file {transfer} {transfer.getFileName()}: {error}') + self.listener.onDownloadError(error.toString()) + self.error = error.toString() + self.continue_event.set() + + def cancel_download(self): + self.is_cancelled = True + +class AsyncExecutor(object): + + def __init__(self): + self.continue_event = threading.Event() + + def do(self, function, args): + self.continue_event.clear() + function(*args) + self.continue_event.wait() + +class MegaDownloadHelper: + def __init__(self): + pass + + def add_download(self, mega_link: str, path: str, listener): + executor = AsyncExecutor() + api = MegaApi(MEGA_API_KEY, None, None, 'telegram-mirror-bot') + mega_listener = MegaAppListener(executor.continue_event, listener) + os.makedirs(path) + api.addListener(mega_listener) + executor.do(api.getPublicNode, (mega_link,)) + node = mega_listener.node + if node is None: + executor.do(api.loginToFolder, (mega_link,)) + node = mega_listener.node + if mega_listener.error is not None: + return listener.onDownloadError(str(mega_listener.error)) + mega_listener.setValues(node.getName(),api.getSize(node),mega_link.split("!",1)[-1].split("!",1)[0]) + with download_dict_lock: + download_dict[listener.uid] = MegaDownloadStatus(mega_listener,listener) + threading.Thread(target=executor.do,args=(api.startDownload,(node,path))).start() + update_all_messages() diff --git a/bot/helper/mirror_utils/status_utils/mega_download_status.py b/bot/helper/mirror_utils/status_utils/mega_download_status.py new file mode 100644 index 0000000..bb5cf73 --- /dev/null +++ b/bot/helper/mirror_utils/status_utils/mega_download_status.py @@ -0,0 +1,62 @@ +from bot.helper.ext_utils.bot_utils import get_readable_file_size,MirrorStatus, get_readable_time +from bot import DOWNLOAD_DIR +from .status import Status + + +class MegaDownloadStatus(Status): + + def __init__(self, obj,listener): + self.uid = obj.uid + self.listener = listener + self.obj = obj + + def name(self) -> str: + return self.obj.name + + def progress_raw(self): + try: + return round(self.processed_bytes() / self.obj.size * 100,2) + except ZeroDivisionError: + return 0.0 + + def progress(self): + """Progress of download in percentage""" + return f"{self.progress_raw()}%" + + + def status(self) -> str: + return MirrorStatus.STATUS_DOWNLOADING + + def processed_bytes(self): + return self.obj.downloaded_bytes + + def eta(self): + try: + seconds = (self.size_raw() - self.processed_bytes()) / self.speed_raw() + return f'{get_readable_time(seconds)}' + except ZeroDivisionError: + return '-' + + def size_raw(self): + return self.obj.size + + def size(self) -> str: + return get_readable_file_size(self.size_raw()) + + def downloaded(self) -> str: + return get_readable_file_size(self.obj.downloadedBytes) + + def speed_raw(self): + return self.obj.speed + + def speed(self) -> str: + return f'{get_readable_file_size(self.speed_raw())}/s' + + def gid(self) -> str: + return self.obj.gid + + def path(self) -> str: + return f"{DOWNLOAD_DIR}{self.uid}" + + def download(self): + return self.obj \ No newline at end of file diff --git a/bot/modules/mirror.py b/bot/modules/mirror.py index b253139..a6d2cd0 100644 --- a/bot/modules/mirror.py +++ b/bot/modules/mirror.py @@ -7,6 +7,7 @@ from bot.helper.ext_utils import fs_utils, bot_utils from bot.helper.ext_utils.bot_utils import setInterval from bot.helper.ext_utils.exceptions import DirectDownloadLinkException from bot.helper.mirror_utils.download_utils.aria2_download import AriaDownloadHelper +from bot.helper.mirror_utils.download_utils.mega_downloader import MegaDownloadHelper from bot.helper.mirror_utils.download_utils.direct_link_generator import direct_link_generator from bot.helper.mirror_utils.download_utils.telegram_downloader import TelegramDownloadHelper from bot.helper.mirror_utils.status_utils import listeners @@ -211,7 +212,11 @@ def _mirror(bot, update, isTar=False, extract=False): except DirectDownloadLinkException as e: LOGGER.info(f'{link}: {e}') listener = MirrorListener(bot, update, isTar, tag, extract) - ariaDlManager.add_download(link, f'{DOWNLOAD_DIR}/{listener.uid}/',listener) + if bot_utils.is_mega_link(link): + mega_dl = MegaDownloadHelper() + mega_dl.add_download(link,f'{DOWNLOAD_DIR}/{listener.uid}/',listener) + else: + ariaDlManager.add_download(link, f'{DOWNLOAD_DIR}/{listener.uid}/',listener) sendStatusMessage(update, bot) if len(Interval) == 0: Interval.append(setInterval(DOWNLOAD_STATUS_UPDATE_INTERVAL, update_all_messages))