diff --git a/bot/helper/ext_utils/bot_utils.py b/bot/helper/ext_utils/bot_utils.py index d26225a4ace..b79f58cef4b 100644 --- a/bot/helper/ext_utils/bot_utils.py +++ b/bot/helper/ext_utils/bot_utils.py @@ -139,6 +139,8 @@ def is_url(url: str): return True return False +def is_gdrive_link(url: str): + return "drive.google.com" in url def is_mega_link(url: str): return "mega.nz" in url diff --git a/bot/helper/mirror_utils/download_utils/direct_link_generator.py b/bot/helper/mirror_utils/download_utils/direct_link_generator.py index b2163b0d530..0eca3e9e904 100644 --- a/bot/helper/mirror_utils/download_utils/direct_link_generator.py +++ b/bot/helper/mirror_utils/download_utils/direct_link_generator.py @@ -34,8 +34,6 @@ def direct_link_generator(link: str): raise DirectDownloadLinkException("`No links found!`") elif 'youtube.com' in link or 'youtu.be' in link: raise DirectDownloadLinkException(f"Youtube Link use /{BotCommands.WatchCommand} or /{BotCommands.TarWatchCommand}") - elif 'drive.google.com' in link: - raise DirectDownloadLinkException(f"G-Drive Link use /{BotCommands.CloneCommand}") elif 'zippyshare.com' in link: return zippy_share(link) elif 'yadi.sk' in link: diff --git a/bot/helper/mirror_utils/status_utils/gdownload_status.py b/bot/helper/mirror_utils/status_utils/gdownload_status.py new file mode 100644 index 00000000000..83f474911ed --- /dev/null +++ b/bot/helper/mirror_utils/status_utils/gdownload_status.py @@ -0,0 +1,61 @@ +from .status import Status +from bot.helper.ext_utils.bot_utils import MirrorStatus, get_readable_file_size, get_readable_time +from bot import DOWNLOAD_DIR + + +class DownloadStatus(Status): + def __init__(self, obj, size, listener, gid): + self.dobj = obj + self.__dsize = size + self.uid = listener.uid + self.message = listener.message + self.__dgid = gid + + def path(self): + return f"{DOWNLOAD_DIR}{self.uid}" + + def processed_bytes(self): + return self.dobj.downloaded_bytes + + def size_raw(self): + return self.__dsize + + def size(self): + return get_readable_file_size(self.__dsize) + + def status(self): + return MirrorStatus.STATUS_DOWNLOADING + + def name(self): + return self.dobj.name + + def gid(self) -> str: + return self.__dgid + + def progress_raw(self): + try: + return self.dobj.downloaded_bytes / self.__dsize * 100 + except ZeroDivisionError: + return 0 + + def progress(self): + return f'{round(self.progress_raw(), 2)}%' + + def speed_raw(self): + """ + :return: Download speed in Bytes/Seconds + """ + return self.dobj.dspeed() + + def speed(self): + return f'{get_readable_file_size(self.speed_raw())}/s' + + def eta(self): + try: + seconds = (self.__dsize - self.dobj.downloaded_bytes) / self.speed_raw() + return f'{get_readable_time(seconds)}' + except ZeroDivisionError: + return '-' + + def download(self): + return self.dobj diff --git a/bot/helper/mirror_utils/upload_utils/gdriveTools.py b/bot/helper/mirror_utils/upload_utils/gdriveTools.py index 99d395c074d..41904710d6b 100644 --- a/bot/helper/mirror_utils/upload_utils/gdriveTools.py +++ b/bot/helper/mirror_utils/upload_utils/gdriveTools.py @@ -1,4 +1,5 @@ import os +import io import pickle import urllib.parse as urlparse from urllib.parse import parse_qs @@ -13,7 +14,7 @@ from google_auth_oauthlib.flow import InstalledAppFlow from googleapiclient.discovery import build from googleapiclient.errors import HttpError -from googleapiclient.http import MediaFileUpload +from googleapiclient.http import MediaFileUpload, MediaIoBaseDownload from tenacity import * from telegram import InlineKeyboardMarkup @@ -43,14 +44,18 @@ def __init__(self, name=None, listener=None): self.__service = self.authorize() self.__listener = listener self._file_uploaded_bytes = 0 + self._file_downloaded_bytes = 0 self.uploaded_bytes = 0 + self.downloaded_bytes = 0 self.UPDATE_INTERVAL = 5 self.start_time = 0 self.total_time = 0 + self.dtotal_time = 0 self._should_update = True self.is_uploading = True self.is_cancelled = False self.status = None + self.dstatus = None self.updater = None self.name = name self.update_interval = 3 @@ -74,6 +79,12 @@ def speed(self): except ZeroDivisionError: return 0 + def dspeed(self): + try: + return self.downloaded_bytes / self.dtotal_time + except ZeroDivisionError: + return 0 + @staticmethod def getIdFromUrl(link: str): if "folders" in link or "file" in link: @@ -728,3 +739,106 @@ def clonehelper(self, link): return msg, "", "" return "", clonesize, name + def download(self, link): + self.is_downloading = True + file_id = self.getIdFromUrl(link) + if USE_SERVICE_ACCOUNTS: + self.service_account_count = len(os.listdir("accounts")) + self.start_time = time.time() + self.updater = setInterval(self.update_interval, self._on_download_progress) + try: + meta = self.getFileMetadata(file_id) + path = f"{DOWNLOAD_DIR}{self.__listener.uid}/" + if meta.get("mimeType") == self.__G_DRIVE_DIR_MIME_TYPE: + self.download_folder(file_id, path, meta.get('name')) + else: + os.makedirs(path) + self.download_file(file_id, path, meta.get('name'), meta.get('mimeType')) + except Exception as err: + if isinstance(err, RetryError): + LOGGER.info(f"Total Attempts: {err.last_attempt.attempt_number}") + err = err.last_attempt.exception() + err = str(err).replace('>', '').replace('<', '') + LOGGER.error(err) + self.is_cancelled = True + self.__listener.onDownloadError(err) + return + finally: + self.updater.cancel() + if self.is_cancelled: + return + self.__listener.onDownloadComplete() + + def download_folder(self, folder_id, path, folder_name): + if not os.path.exists(path + folder_name): + os.makedirs(path + folder_name) + path += folder_name + '/' + result = [] + page_token = None + while True: + files = self.__service.files().list( + supportsTeamDrives=True, + includeTeamDriveItems=True, + q=f"'{folder_id}' in parents", + fields='nextPageToken, files(id, name, mimeType, size, shortcutDetails)', + pageToken=page_token, + pageSize=1000).execute() + result.extend(files['files']) + page_token = files.get("nextPageToken") + if not page_token: + break + + result = sorted(result, key=lambda k: k['name']) + for item in result: + file_id = item['id'] + filename = item['name'] + mime_type = item['mimeType'] + shortcut_details = item.get('shortcutDetails', None) + if shortcut_details != None: + file_id = shortcut_details['targetId'] + mime_type = shortcut_details['targetMimeType'] + if mime_type == 'application/vnd.google-apps.folder': + self.download_folder(file_id, path, filename) + elif not os.path.isfile(path + filename): + self.download_file(file_id, path, filename, mime_type) + if self.is_cancelled: + break + return + + def download_file(self, file_id, path, filename, mime_type): + request = self.__service.files().get_media(fileId=file_id) + fh = io.FileIO('{}{}'.format(path, filename), 'wb') + downloader = MediaIoBaseDownload(fh, request, chunksize = 50 * 1024 * 1024) + done = False + while done is False: + if self.is_cancelled: + fh.close() + break + return + try: + self.dstatus, done = downloader.next_chunk() + except HttpError as err: + if err.resp.get('content-type', '').startswith('application/json'): + reason = json.loads(err.content).get('error').get('errors')[0].get('reason') + if reason == 'userRateLimitExceeded' or reason == 'dailyLimitExceeded': + if USE_SERVICE_ACCOUNTS: + if not self.switchServiceAccount(): + raise err + LOGGER.info(f"Got: {reason}, Trying Again...") + return self.download_file(file_id, path, filename, mime_type) + else: + raise err + else: + raise err + self._file_downloaded_bytes = 0 + + def _on_download_progress(self): + if self.dstatus is not None: + chunk_size = self.dstatus.total_size * self.dstatus.progress() - self._file_downloaded_bytes + self._file_downloaded_bytes = self.dstatus.total_size * self.dstatus.progress() + self.downloaded_bytes += chunk_size + self.dtotal_time += self.update_interval + + def cancel_download(self): + self.is_cancelled = True + self.__listener.onDownloadError('Download stopped by user!') diff --git a/bot/helper/telegram_helper/message_utils.py b/bot/helper/telegram_helper/message_utils.py index 36a0d3d0e2f..e1f120574e9 100644 --- a/bot/helper/telegram_helper/message_utils.py +++ b/bot/helper/telegram_helper/message_utils.py @@ -80,9 +80,9 @@ def update_all_messages(): for download in list(download_dict.values()): speedy = download.speed() if download.status() == MirrorStatus.STATUS_DOWNLOADING: - if 'KiB/s' in speedy: + if 'K' in speedy: dlspeed_bytes += float(speedy.split('K')[0]) * 1024 - elif 'MiB/s' in speedy: + elif 'M' in speedy: dlspeed_bytes += float(speedy.split('M')[0]) * 1048576 if download.status() == MirrorStatus.STATUS_UPLOADING: if 'KB/s' in speedy: @@ -118,9 +118,9 @@ def sendStatusMessage(msg, bot): for download in list(download_dict.values()): speedy = download.speed() if download.status() == MirrorStatus.STATUS_DOWNLOADING: - if 'KiB/s' in speedy: + if 'K' in speedy: dlspeed_bytes += float(speedy.split('K')[0]) * 1024 - elif 'MiB/s' in speedy: + elif 'M' in speedy: dlspeed_bytes += float(speedy.split('M')[0]) * 1048576 if download.status() == MirrorStatus.STATUS_UPLOADING: if 'KB/s' in speedy: diff --git a/bot/modules/mirror.py b/bot/modules/mirror.py index 422558bedbe..b99c8db948a 100644 --- a/bot/modules/mirror.py +++ b/bot/modules/mirror.py @@ -15,6 +15,7 @@ from bot.helper.mirror_utils.status_utils.extract_status import ExtractStatus from bot.helper.mirror_utils.status_utils.tar_status import TarStatus from bot.helper.mirror_utils.status_utils.upload_status import UploadStatus +from bot.helper.mirror_utils.status_utils.gdownload_status import DownloadStatus from bot.helper.mirror_utils.upload_utils import gdriveTools from bot.helper.telegram_helper.bot_commands import BotCommands from bot.helper.telegram_helper.filters import CustomFilters @@ -26,6 +27,8 @@ import subprocess import threading import re +import random +import string ariaDlManager = AriaDownloadHelper() ariaDlManager.start_listener() @@ -291,14 +294,32 @@ def _mirror(bot, update, isTar=False, extract=False): if "ERROR:" in str(e): sendMessage(f"{e}", bot, update) return - if "G-Drive" in str(e): - sendMessage(f"ERROR: {e}", bot, update) - return if "Youtube" in str(e): sendMessage(f"ERROR: {e}", bot, update) return + listener = MirrorListener(bot, update, pswd, isTar, tag, extract) - if bot_utils.is_mega_link(link): + + if bot_utils.is_gdrive_link(link): + if not isTar and not extract: + sendMessage(f"Use /{BotCommands.CloneCommand} To Copy File/Folder", bot, update) + return + res, size, name = gdriveTools.GoogleDriveHelper().clonehelper(link) + if res != "": + sendMessage(res, bot, update) + return + LOGGER.info(f"Download Name : {name}") + drive = gdriveTools.GoogleDriveHelper(name, listener) + gid = ''.join(random.SystemRandom().choices(string.ascii_letters + string.digits, k=12)) + download_status = DownloadStatus(drive, size, listener, gid) + with download_dict_lock: + download_dict[listener.uid] = download_status + if len(Interval) == 0: + Interval.append(setInterval(DOWNLOAD_STATUS_UPDATE_INTERVAL, update_all_messages)) + sendStatusMessage(update, bot) + drive.download(link) + + elif bot_utils.is_mega_link(link): link_type = get_mega_link_type(link) if link_type == "folder" and BLOCK_MEGA_FOLDER: sendMessage("Mega folder are blocked!", bot, update)