Mega: Fix folder downloads with mega account

bot/helper/ext_utils/bot_utils.py

@@ -143,6 +143,15 @@ def is_magnet(url: str):
 def is_mega_link(url: str):
     return "mega.nz" in url
 
+def get_mega_link_type(url: str):
+    if "folder" in url:
+        return "folder"
+    elif "file" in url:
+        return "file"
+    elif "/#F!" in url:
+        return "folder"
+    return "file"
+
 
 def new_thread(fn):
     """To use as decorator to make a function call threaded.

bot/helper/mirror_utils/download_utils/mega_downloader.py

@@ -3,6 +3,7 @@
 from mega import (MegaApi, MegaListener, MegaRequest, MegaTransfer, MegaError)
 from bot.helper.telegram_helper.message_utils import update_all_messages
 import os
+from bot.helper.ext_utils.bot_utils import new_thread, get_mega_link_type
 from bot.helper.mirror_utils.status_utils.mega_download_status import MegaDownloadStatus
 import random
 import string
@@ -12,12 +13,13 @@ class MegaDownloaderException(Exception):
 
 
 class MegaAppListener(MegaListener):
-    _NO_EVENT_ON = (MegaRequest.TYPE_LOGIN,
-                    MegaRequest.TYPE_FETCH_NODES)
+    _NO_EVENT_ON = (MegaRequest.TYPE_LOGIN,MegaRequest.TYPE_FETCH_NODES)
+    NO_ERROR = "no error"
 
     def __init__(self, continue_event: threading.Event, listener):
         self.continue_event = continue_event
         self.node = None
+        self.public_node = None
         self.listener = listener
         self.uid = listener.uid
         self.__bytes_transferred = 0
@@ -64,15 +66,19 @@ def onRequestFinish(self, api, request, error):
         if request_type == MegaRequest.TYPE_LOGIN:
             api.fetchNodes()
         elif request_type == MegaRequest.TYPE_GET_PUBLIC_NODE:
-            self.node = request.getPublicMegaNode()
+            self.public_node = request.getPublicMegaNode()
         elif request_type == MegaRequest.TYPE_FETCH_NODES:
             LOGGER.info("Fetching Root Node.")
             self.node = api.getRootNode()
-        if request_type not in self._NO_EVENT_ON:
+            LOGGER.info(f"Node Name: {self.node.getName()}")
+        if request_type not in self._NO_EVENT_ON or self.node and "cloud drive" not in self.node.getName().lower():
             self.continue_event.set()
 
     def onRequestTemporaryError(self, api, request, error: MegaError):
-        self.listener.onDownloadError(error.toString())
+        LOGGER.info(f'Mega Request error in {error}')
+        if not self.is_cancelled:
+            self.listener.onDownloadError("RequestTempError: " + error.toString())
+            self.is_cancelled = True
         self.error = error.toString()
         self.continue_event.set()
 
@@ -88,22 +94,22 @@ def onTransferUpdate(self, api: MegaApi, transfer: MegaTransfer):
     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:
+            if transfer.isFolderTransfer() and transfer.isFinished() or transfer.getFileName() == self.name and not self.is_cancelled:
                 self.listener.onDownloadComplete()
+                self.continue_event.set()
         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()
+        if not self.is_cancelled:
+            self.is_cancelled = True
+            self.listener.onDownloadError("TransferTempError: "+self.error)
 
     def cancel_download(self):
         self.is_cancelled = True
+        self.listener.onDownloadError("Download Canceled by user")
 
 
 class AsyncExecutor:
@@ -122,26 +128,31 @@ def __init__(self):
         pass
 
     @staticmethod
+    @new_thread
     def add_download(mega_link: str, path: str, listener):
         if MEGA_API_KEY is None:
             raise MegaDownloaderException('Mega API KEY not provided! Cannot mirror mega links')
         executor = AsyncExecutor()
         api = MegaApi(MEGA_API_KEY, None, None, 'telegram-mirror-bot')
         mega_listener = MegaAppListener(executor.continue_event, listener)
+        with download_dict_lock:
+            download_dict[listener.uid] = MegaDownloadStatus(mega_listener, listener)
         os.makedirs(path)
         api.addListener(mega_listener)
         if MEGA_EMAIL_ID is not None and MEGA_PASSWORD is not None:
             executor.do(api.login, (MEGA_EMAIL_ID, MEGA_PASSWORD))
-        executor.do(api.getPublicNode, (mega_link,))
-        node = mega_listener.node
-        if node is None:
-            executor.do(api.loginToFolder, (mega_link,))
-            node = mega_listener.node
+        link_type = get_mega_link_type(mega_link)
+        if link_type == "file":
+            executor.do(api.getPublicNode, (mega_link,))
+            node = mega_listener.public_node
+        else:
+            LOGGER.info("Logging into mega folder")
+            folder_api = MegaApi(MEGA_API_KEY,None,None,'TgBot')
+            folder_api.addListener(mega_listener)
+            executor.do(folder_api.loginToFolder, (mega_link,))
+            node = folder_api.authorizeNode(mega_listener.node)
         if mega_listener.error is not None:
             return listener.onDownloadError(str(mega_listener.error))
         gid = ''.join(random.SystemRandom().choices(string.ascii_letters + string.digits, k=8))
         mega_listener.setValues(node.getName(), api.getSize(node), gid)
-        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()
+        executor.do(api.startDownload,(node,path))