Incomplete mega support

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)

bot/helper/ext_utils/bot_utils.py

@@ -61,14 +61,22 @@ 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:
             return index
         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:

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
+

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:

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}"

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'

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