tar/unzip G-Drive Links (#220)

* tar/unzip G-Drive Links

Signed-off-by: anas <[email protected]>

* Fix overall download speed for mega and others

Signed-off-by: anas <[email protected]>

Author: anasty17

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

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:

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

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

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:

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)