From 5a7814e76c2e90af82fcfea501db9c076ed2021e Mon Sep 17 00:00:00 2001 From: lzzy12 Date: Sun, 27 Oct 2019 08:44:29 +0530 Subject: [PATCH] Incomplete mega support Signed-off-by: lzzy12 --- bot/__init__.py | 1 + bot/helper/ext_utils/bot_utils.py | 10 +- bot/helper/ext_utils/exceptions.py | 7 ++ bot/helper/mirror_utils/download_tools.py | 8 +- bot/helper/mirror_utils/gdriveTools.py | 4 +- bot/helper/mirror_utils/mega_status.py | 39 +++++++ bot/helper/mirror_utils/mega_tools.py | 127 ++++++++++++++++++++++ 7 files changed, 190 insertions(+), 6 deletions(-) create mode 100644 bot/helper/mirror_utils/mega_status.py create mode 100644 bot/helper/mirror_utils/mega_tools.py diff --git a/bot/__init__.py b/bot/__init__.py index 8636ec9..ccec067 100644 --- a/bot/__init__.py +++ b/bot/__init__.py @@ -65,6 +65,7 @@ def getConfig(name: str): DOWNLOAD_STATUS_UPDATE_INTERVAL = int(getConfig('DOWNLOAD_STATUS_UPDATE_INTERVAL')) OWNER_ID = int(getConfig('OWNER_ID')) AUTO_DELETE_MESSAGE_DURATION = int(getConfig('AUTO_DELETE_MESSAGE_DURATION')) + MEGA_API_KEY = getConfig('MEGA_API_KEY') except KeyError as e: LOGGER.error("One or more env variables missing! Exiting now") exit(1) diff --git a/bot/helper/ext_utils/bot_utils.py b/bot/helper/ext_utils/bot_utils.py index 2d7553b..a006cbf 100644 --- a/bot/helper/ext_utils/bot_utils.py +++ b/bot/helper/ext_utils/bot_utils.py @@ -61,7 +61,7 @@ def get_progress_bar_string(status): return p_str -def get_download_index(_list, gid): +def get_download_index_by_aria_gid(_list, gid): index = 0 for i in _list: if i.download().gid == gid: @@ -69,6 +69,14 @@ def get_download_index(_list, gid): index += 1 +def get_download_index_by_uid(_list, uid): + index = 0 + for i in _list: + if i.uid() == uid: + return index + index += 1 + + def get_download_str(): result = "" with download_dict_lock: diff --git a/bot/helper/ext_utils/exceptions.py b/bot/helper/ext_utils/exceptions.py index e39e15a..33d06b1 100644 --- a/bot/helper/ext_utils/exceptions.py +++ b/bot/helper/ext_utils/exceptions.py @@ -15,3 +15,10 @@ class DownloadCancelled(Exception): def __init__(self, message, error=None): super().__init__(message) self.error = error + + +class MegaDownloadError(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 237d87f..7afb131 100644 --- a/bot/helper/mirror_utils/download_tools.py +++ b/bot/helper/mirror_utils/download_tools.py @@ -1,10 +1,10 @@ +from bot.helper.mirror_utils.mega_tools import mega_download from time import sleep from bot import DOWNLOAD_DIR, DOWNLOAD_STATUS_UPDATE_INTERVAL, aria2 from .download_status import DownloadStatus from bot.helper.ext_utils.bot_utils import * from bot.helper.ext_utils.exceptions import KillThreadException - class DownloadHelper: def __init__(self, listener=None): @@ -13,6 +13,8 @@ def __init__(self, listener=None): def add_download(self, link: str): if is_url(link): + if 'mega.nz' in link: + mega_download(link, DOWNLOAD_DIR + str(self.__listener.uid), self.__listener) if link.endswith('.torrent'): self.__is_torrent = True download = aria2.add_uris([link], {'dir': DOWNLOAD_DIR + str(self.__listener.uid)}) @@ -36,7 +38,7 @@ def __get_followed_download_gid(self): def __update_download_status(self): status_list = get_download_status_list() - index = get_download_index(status_list, self.__get_download().gid) + index = get_download_index_by_aria_gid(status_list, self.__get_download().gid) # This tracks if message exists or did it get replaced by other status message should_update = True if self.__is_torrent: @@ -74,7 +76,7 @@ def __update_download_status(self): return if should_update: status_list = get_download_status_list() - index = get_download_index(status_list, self.__get_download().gid) + index = get_download_index_by_aria_gid(status_list, self.__get_download().gid) # TODO: Find a better way to differentiate between 2 list of objects progress_str_list = get_download_str() if progress_str_list != previous: diff --git a/bot/helper/mirror_utils/gdriveTools.py b/bot/helper/mirror_utils/gdriveTools.py index b96a75d..f862f89 100644 --- a/bot/helper/mirror_utils/gdriveTools.py +++ b/bot/helper/mirror_utils/gdriveTools.py @@ -61,7 +61,7 @@ def _on_upload_progress(self): try: LOGGER.info('Updating messages') _list = get_download_status_list() - index = get_download_index(_list, get_download(self.__listener.message.message_id).gid) + index = get_download_index_by_aria_gid(_list, get_download(self.__listener.message.message_id).gid) self.__listener.onUploadProgress(_list, index) except KillThreadException as e: LOGGER.info(f'Stopped calling onDownloadProgress(): {str(e)}') @@ -110,7 +110,7 @@ def upload_file(self, file_path, file_name, mime_type, parent_id): def upload(self, file_name: str): _list = get_download_status_list() - index = get_download_index(_list, get_download(self.__listener.message.message_id).gid) + index = get_download_index_by_aria_gid(_list, get_download(self.__listener.message.message_id).gid) self.__listener.onUploadStarted(_list, index) file_dir = f"{DOWNLOAD_DIR}{self.__listener.message.message_id}" file_path = f"{file_dir}/{file_name}" diff --git a/bot/helper/mirror_utils/mega_status.py b/bot/helper/mirror_utils/mega_status.py new file mode 100644 index 0000000..52307ea --- /dev/null +++ b/bot/helper/mirror_utils/mega_status.py @@ -0,0 +1,39 @@ +from bot.helper.ext_utils.bot_utils import get_readable_file_size + + +class MegaDownloadStatus: + STATUS_UPLOADING = "Uploading" + STATUS_DOWNLOADING = "Downloading" + STATUS_WAITING = "Queued" + STATUS_FAILED = "Failed. Cleaning download" + STATUS_CANCELLED = "Cancelled" + + def __init__(self, uid): + self.uid = uid + self.name = '' + self.downloadedBytes = 0 + self.sizeBytes = 0 + self.speed = 0 + self.status = '' + + def name(self) -> str: + return self.name + + def progress(self): + """Progress of download in percentage""" + return (self.downloadedBytes // self.sizeBytes) * 100 + + def status(self) -> str: + return self.status + + def eta(self) -> str: + return get_readable_file_size(self.sizeBytes / self.speed) + + def size(self) -> str: + return get_readable_file_size(self.size) + + def downloaded(self) -> str: + return get_readable_file_size(self.downloadedBytes) + + def speed(self) -> str: + return f'{get_readable_file_size(self.downloadedBytes)}/s' \ No newline at end of file diff --git a/bot/helper/mirror_utils/mega_tools.py b/bot/helper/mirror_utils/mega_tools.py new file mode 100644 index 0000000..c5b9545 --- /dev/null +++ b/bot/helper/mirror_utils/mega_tools.py @@ -0,0 +1,127 @@ +from bot import LOGGER, MEGA_API_KEY +import threading +from mega import (MegaApi, MegaListener, MegaRequest, MegaTransfer, MegaError) +from bot.helper.mirror_utils.listeners import MirrorListeners +from bot.helper.ext_utils.bot_utils import * +import os + + +class MegaAppListener(MegaListener): + + def __init__(self, continue_event: threading.Event, listener: MirrorListeners): + self.continue_event = continue_event + self.node = None + self.listener = listener + self.__bytes_transferred = 0 + self.__speed = 0 + self.__name = '' + self.__size = 0 + self.__should_update = True + self.error = None + 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 + + @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: + api.fetchNodes() + if request_type == MegaRequest.TYPE_FETCH_NODES: + self.node = api.getRootNode() + if request_type != MegaRequest.TYPE_LOGIN: + self.continue_event.set() + + def onRequestTemporaryError(self, api, request, error: MegaError): + _list = get_download_status_list() + _index = get_download_index_by_uid(_list, self.listener.uid) + self.listener.onDownloadError(error.toString(), _list, _index) + self.error = error.toString() + self.continue_event.set() + + def onTransferStart(self, api: MegaApi, transfer: MegaTransfer): + self.__name = transfer.getFileName() + self.__size = transfer.getTotalBytes() + + def onTransferUpdate(self, api: MegaApi, transfer: MegaTransfer): + self.__speed = transfer.getSpeed() + self.__bytes_transferred = transfer.getTransferredBytes() + _list = get_download_status_list() + self.listener.onDownloadProgress(_list, get_download_index_by_uid(_list, self.listener.uid)) + + def onTransferFinish(self, api, transfer, error): + try: + LOGGER.info(f'Transfer finished ({transfer}); Result: {transfer.getFileName()}') + _list = get_download_status_list() + index = get_download_index_by_uid(_list, self.listener.uid) + self.listener.onDownloadComplete(_list, index) + except Exception as e: + LOGGER.error(e) + finally: + self.continue_event.set() + + def onTransferTemporaryError(self, api, transfer, error): + LOGGER.info(f'Mega download error in file {transfer} {transfer.getFileName()}: {error}') + _list = get_download_status_list() + _index = get_download_index_by_uid(_list, self.listener.uid) + self.listener.onDownloadError(error.toString(), _list, _index) + self.error = error.toString() + self.continue_event.set() + + +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 mega_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 not path.endswith('/') or path.endswith('\\'): + path += '/' + + if node is None: + executor.do(api.loginToFolder, (mega_link,)) + if mega_listener.error is not None: + return + listener.onDownloadStarted(mega_link) + executor.do(api.startDownload, (node, path))