Implement Mega.nz mirror support.

- Co-authored by - lzzy12 (github.com/lzzy12)

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

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()

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

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))