|
| 1 | +from bot import LOGGER, MEGA_API_KEY, download_dict_lock, download_dict, MEGA_EMAIL_ID, MEGA_PASSWORD |
| 2 | +import threading |
| 3 | +from mega import (MegaApi, MegaListener, MegaRequest, MegaTransfer, MegaError) |
| 4 | +from bot.helper.telegram_helper.message_utils import update_all_messages |
| 5 | +import os |
| 6 | +from bot.helper.mirror_utils.status_utils.mega_download_status import MegaDownloadStatus |
| 7 | +import random |
| 8 | +import string |
| 9 | + |
| 10 | +class MegaDownloaderException(Exception): |
| 11 | + pass |
| 12 | + |
| 13 | + |
| 14 | +class MegaAppListener(MegaListener): |
| 15 | + _NO_EVENT_ON = (MegaRequest.TYPE_LOGIN, |
| 16 | + MegaRequest.TYPE_FETCH_NODES) |
| 17 | + |
| 18 | + def __init__(self, continue_event: threading.Event, listener): |
| 19 | + self.continue_event = continue_event |
| 20 | + self.node = None |
| 21 | + self.listener = listener |
| 22 | + self.uid = listener.uid |
| 23 | + self.__bytes_transferred = 0 |
| 24 | + self.is_cancelled = False |
| 25 | + self.__speed = 0 |
| 26 | + self.__name = '' |
| 27 | + self.__size = 0 |
| 28 | + self.error = None |
| 29 | + self.gid = "" |
| 30 | + super(MegaAppListener, self).__init__() |
| 31 | + |
| 32 | + @property |
| 33 | + def speed(self): |
| 34 | + """Returns speed of the download in bytes/second""" |
| 35 | + return self.__speed |
| 36 | + |
| 37 | + @property |
| 38 | + def name(self): |
| 39 | + """Returns name of the download""" |
| 40 | + return self.__name |
| 41 | + |
| 42 | + def setValues(self, name, size, gid): |
| 43 | + self.__name = name |
| 44 | + self.__size = size |
| 45 | + self.gid = gid |
| 46 | + |
| 47 | + @property |
| 48 | + def size(self): |
| 49 | + """Size of download in bytes""" |
| 50 | + return self.__size |
| 51 | + |
| 52 | + @property |
| 53 | + def downloaded_bytes(self): |
| 54 | + return self.__bytes_transferred |
| 55 | + |
| 56 | + def onRequestStart(self, api, request): |
| 57 | + LOGGER.info('Request start ({})'.format(request)) |
| 58 | + |
| 59 | + def onRequestFinish(self, api, request, error): |
| 60 | + LOGGER.info('Mega Request finished ({}); Result: {}' |
| 61 | + .format(request, error)) |
| 62 | + |
| 63 | + request_type = request.getType() |
| 64 | + if request_type == MegaRequest.TYPE_LOGIN: |
| 65 | + api.fetchNodes() |
| 66 | + elif request_type == MegaRequest.TYPE_GET_PUBLIC_NODE: |
| 67 | + self.node = request.getPublicMegaNode() |
| 68 | + elif request_type == MegaRequest.TYPE_FETCH_NODES: |
| 69 | + LOGGER.info("Fetching Root Node.") |
| 70 | + self.node = api.getRootNode() |
| 71 | + if request_type not in self._NO_EVENT_ON: |
| 72 | + self.continue_event.set() |
| 73 | + |
| 74 | + def onRequestTemporaryError(self, api, request, error: MegaError): |
| 75 | + self.listener.onDownloadError(error.toString()) |
| 76 | + self.error = error.toString() |
| 77 | + self.continue_event.set() |
| 78 | + |
| 79 | + def onTransferStart(self, api: MegaApi, transfer: MegaTransfer): |
| 80 | + LOGGER.info(f"Transfer Started: {transfer.getFileName()}") |
| 81 | + |
| 82 | + def onTransferUpdate(self, api: MegaApi, transfer: MegaTransfer): |
| 83 | + if self.is_cancelled: |
| 84 | + api.cancelTransfer(transfer, None) |
| 85 | + self.__speed = transfer.getSpeed() |
| 86 | + self.__bytes_transferred = transfer.getTransferredBytes() |
| 87 | + |
| 88 | + def onTransferFinish(self, api: MegaApi, transfer: MegaTransfer, error): |
| 89 | + try: |
| 90 | + LOGGER.info(f'Transfer finished ({transfer}); Result: {transfer.getFileName()}') |
| 91 | + if str(error) != "No error" and self.is_cancelled: |
| 92 | + self.is_cancelled = False |
| 93 | + return self.listener.onDownloadError(error.toString()) |
| 94 | + if transfer.isFolderTransfer() and transfer.isFinished() and not self.is_cancelled or transfer.getFileName() == self.name and not self.is_cancelled: |
| 95 | + self.listener.onDownloadComplete() |
| 96 | + except Exception as e: |
| 97 | + LOGGER.error(e) |
| 98 | + |
| 99 | + def onTransferTemporaryError(self, api, transfer, error): |
| 100 | + LOGGER.info(f'Mega download error in file {transfer} {transfer.getFileName()}: {error}') |
| 101 | + self.listener.onDownloadError(error.toString()) |
| 102 | + self.error = error.toString() |
| 103 | + self.continue_event.set() |
| 104 | + |
| 105 | + def cancel_download(self): |
| 106 | + self.is_cancelled = True |
| 107 | + |
| 108 | + |
| 109 | +class AsyncExecutor: |
| 110 | + |
| 111 | + def __init__(self): |
| 112 | + self.continue_event = threading.Event() |
| 113 | + |
| 114 | + def do(self, function, args): |
| 115 | + self.continue_event.clear() |
| 116 | + function(*args) |
| 117 | + self.continue_event.wait() |
| 118 | + |
| 119 | + |
| 120 | +class MegaDownloadHelper: |
| 121 | + def __init__(self): |
| 122 | + pass |
| 123 | + |
| 124 | + @staticmethod |
| 125 | + def add_download(mega_link: str, path: str, listener): |
| 126 | + if MEGA_API_KEY is None: |
| 127 | + raise MegaDownloaderException('Mega API KEY not provided! Cannot mirror mega links') |
| 128 | + executor = AsyncExecutor() |
| 129 | + api = MegaApi(MEGA_API_KEY, None, None, 'telegram-mirror-bot') |
| 130 | + mega_listener = MegaAppListener(executor.continue_event, listener) |
| 131 | + os.makedirs(path) |
| 132 | + api.addListener(mega_listener) |
| 133 | + if MEGA_EMAIL_ID is not None and MEGA_PASSWORD is not None: |
| 134 | + executor.do(api.login, (MEGA_EMAIL_ID, MEGA_PASSWORD)) |
| 135 | + executor.do(api.getPublicNode, (mega_link,)) |
| 136 | + node = mega_listener.node |
| 137 | + if node is None: |
| 138 | + executor.do(api.loginToFolder, (mega_link,)) |
| 139 | + node = mega_listener.node |
| 140 | + if mega_listener.error is not None: |
| 141 | + return listener.onDownloadError(str(mega_listener.error)) |
| 142 | + gid = ''.join(random.SystemRandom().choices(string.ascii_letters + string.digits, k=8)) |
| 143 | + mega_listener.setValues(node.getName(), api.getSize(node), gid) |
| 144 | + with download_dict_lock: |
| 145 | + download_dict[listener.uid] = MegaDownloadStatus(mega_listener, listener) |
| 146 | + threading.Thread(target=executor.do, args=(api.startDownload, (node, path))).start() |
| 147 | + update_all_messages() |
0 commit comments