From 0ed1ca3531ac942b2352f7977ab5372c118f7e70 Mon Sep 17 00:00:00 2001 From: lzzy12 Date: Tue, 24 Dec 2019 17:52:17 +0530 Subject: [PATCH] Enable TeamDrive support Signed-off-by: lzzy12 --- README.md | 1 + bot/__init__.py | 4 ++++ .../mirror_utils/upload_utils/gdriveTools.py | 20 +++++++++++-------- bot/modules/mirror.py | 8 ++++++-- config_sample.env | 3 ++- 5 files changed, 25 insertions(+), 11 deletions(-) diff --git a/README.md b/README.md index 8ecdfa3..c93794a 100644 --- a/README.md +++ b/README.md @@ -51,6 +51,7 @@ Fill up rest of the fields. Meaning of each fields are discussed below: - DOWNLOAD_STATUS_UPDATE_INTERVAL : A short interval of time in seconds after which the Mirror progress message is updated. (I recommend to keep it 5 seconds at least) - OWNER_ID : The Telegram user ID (not username) of the owner of the bot - AUTO_DELETE_MESSAGE_DURATION : Interval of time (in seconds), after which the bot deletes it's message (and command message) which is expected to be viewed instantly. Note: Set to -1 to never automatically delete messages +- INDEX_URL : (Optional field) Refer to https://github.com/maple3142/GDIndex/ The URL should not have any trailing '/' Note: You can limit maximum concurrent downloads by changing the value of MAX_CONCURRENT_DOWNLOADS in aria.sh. By default, it's set to 2 diff --git a/bot/__init__.py b/bot/__init__.py index 0bd892a..5f03c8f 100644 --- a/bot/__init__.py +++ b/bot/__init__.py @@ -72,6 +72,10 @@ try: except KeyError as e: LOGGER.error("One or more env variables missing! Exiting now") exit(1) +try: + INDEX_URL = getConfig('INDEX_URL') +except KeyError: + INDEX_URL = None updater = tg.Updater(token=BOT_TOKEN) bot = updater.bot diff --git a/bot/helper/mirror_utils/upload_utils/gdriveTools.py b/bot/helper/mirror_utils/upload_utils/gdriveTools.py index bc81b3a..0e39b66 100644 --- a/bot/helper/mirror_utils/upload_utils/gdriveTools.py +++ b/bot/helper/mirror_utils/upload_utils/gdriveTools.py @@ -67,7 +67,8 @@ class GoogleDriveHelper: } if parent_id is not None: file_metadata['parents'] = [parent_id] - return self.__service.files().create(body=file_metadata, media_body=media_body).execute() + return self.__service.files().create(supportsTeamDrives=True, + body=file_metadata, media_body=media_body).execute() def __set_permission(self, drive_id): permissions = { @@ -76,7 +77,7 @@ class GoogleDriveHelper: 'value': None, 'withLink': True } - return self.__service.permissions().create(fileId=drive_id, body=permissions).execute() + return self.__service.permissions().create(supportsTeamDrives=True, fileId=drive_id, body=permissions).execute() def upload_file(self, file_path, file_name, mime_type, parent_id): # File body description @@ -92,7 +93,8 @@ class GoogleDriveHelper: media_body = MediaFileUpload(file_path, mimetype=mime_type, resumable=False) - response = self.__service.files().create(body=file_metadata, media_body=media_body).execute() + response = self.__service.files().create(supportsTeamDrives=True, + body=file_metadata, media_body=media_body).execute() self.__set_permission(response['id']) drive_file = self.__service.files().get(fileId=response['id']).execute() download_url = self.__G_DRIVE_BASE_DOWNLOAD_URL.format(drive_file.get('id')) @@ -103,7 +105,8 @@ class GoogleDriveHelper: chunksize=50*1024*1024) # Insert a file - drive_file = self.__service.files().create(body=file_metadata, media_body=media_body) + drive_file = self.__service.files().create(supportsTeamDrives=True, + body=file_metadata, media_body=media_body) response = None while response is None: if self.is_cancelled: @@ -113,7 +116,7 @@ class GoogleDriveHelper: # Insert new permissions self.__set_permission(response['id']) # Define file instance and get url for download - drive_file = self.__service.files().get(fileId=response['id']).execute() + drive_file = self.__service.files().get(supportsTeamDrives=True, fileId=response['id']).execute() download_url = self.__G_DRIVE_BASE_DOWNLOAD_URL.format(drive_file.get('id')) return download_url @@ -123,7 +126,7 @@ class GoogleDriveHelper: file_path = f"{file_dir}/{file_name}" LOGGER.info("Uploading File: " + file_path) self.start_time = time.time() - self.updater = setInterval(1, self._on_upload_progress) + self.updater = setInterval(5, self._on_upload_progress) if os.path.isfile(file_path): try: mime_type = get_mime_type(file_path) @@ -167,7 +170,7 @@ class GoogleDriveHelper: } if parent_id is not None: file_metadata["parents"] = [parent_id] - file = self.__service.files().create(body=file_metadata).execute() + file = self.__service.files().create(supportsTeamDrives=True, body=file_metadata).execute() file_id = file.get("id") self.__set_permission(file_id) LOGGER.info("Created Google-Drive Folder:\nName: {}\nID: {} ".format(file.get("name"), file_id)) @@ -220,7 +223,8 @@ class GoogleDriveHelper: page_token = None results = [] while True: - response = self.__service.files().list(q=query, + response = self.__service.files().list(supportsTeamDrives=True, + q=query, spaces='drive', fields='nextPageToken, files(id, name, mimeType, size)', pageToken=page_token, diff --git a/bot/modules/mirror.py b/bot/modules/mirror.py index 4f6193a..74aa771 100644 --- a/bot/modules/mirror.py +++ b/bot/modules/mirror.py @@ -6,7 +6,7 @@ from bot.helper.mirror_utils.status_utils.upload_status import UploadStatus from bot.helper.mirror_utils.status_utils.tar_status import TarStatus from bot import dispatcher, DOWNLOAD_DIR, DOWNLOAD_STATUS_UPDATE_INTERVAL from bot.helper.ext_utils import fs_utils, bot_utils -from bot import Interval +from bot import Interval, INDEX_URL from bot.helper.telegram_helper.message_utils import * from bot.helper.ext_utils.bot_utils import setInterval from bot.helper.telegram_helper.filters import CustomFilters @@ -52,7 +52,6 @@ class MirrorListener(listeners.MirrorListeners): path = f'{DOWNLOAD_DIR}{self.uid}/{download_dict[self.uid].name()}' name = pathlib.PurePath(path).name with download_dict_lock: - download_dict[self.uid].is_archiving = False LOGGER.info(f"Upload Name : {name}") drive = gdriveTools.GoogleDriveHelper(name, self) upload_status = UploadStatus(drive, size, self.uid) @@ -92,6 +91,11 @@ class MirrorListener(listeners.MirrorListeners): msg = f'{download_dict[self.uid].name()} ({download_dict[self.uid].size()})' LOGGER.info(f'Done Uploading {download_dict[self.uid].name()}') + share_url = f'{INDEX_URL}/{download_dict[self.uid].name()}' + share_url = share_url.replace(' ', '%20') + if INDEX_URL is not None: + msg += f'\n\n Shareable link: here' + if len(download_dict) == 0: self.clean() else: diff --git a/config_sample.env b/config_sample.env index 92845b8..11aa020 100644 --- a/config_sample.env +++ b/config_sample.env @@ -7,4 +7,5 @@ GDRIVE_FOLDER_ID = "" OWNER_ID = DOWNLOAD_DIR = "/home/username/mirror-bot/downloads" DOWNLOAD_STATUS_UPDATE_INTERVAL = 5 -AUTO_DELETE_MESSAGE_DURATION = 20 \ No newline at end of file +AUTO_DELETE_MESSAGE_DURATION = 20 +INDEX_URL = "" \ No newline at end of file