From edb6c20141cf59ce8c686baa22da6c1791e80cf2 Mon Sep 17 00:00:00 2001 From: Shivam Jha Date: Fri, 10 Jan 2020 17:36:51 +0530 Subject: [PATCH 01/37] Fixed error in gdrivetools --- bot/helper/mirror_utils/upload_utils/gdriveTools.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bot/helper/mirror_utils/upload_utils/gdriveTools.py b/bot/helper/mirror_utils/upload_utils/gdriveTools.py index 059a4f7..7f8f2d7 100644 --- a/bot/helper/mirror_utils/upload_utils/gdriveTools.py +++ b/bot/helper/mirror_utils/upload_utils/gdriveTools.py @@ -7,7 +7,7 @@ import os from tenacity import * import threading -from bot import LOGGER, parent_id, DOWNLOAD_DIR, IS_TEAM_DRIVE, INDEX_URL +from bot import LOGGER, parent_id, DOWNLOAD_DIR, IS_TEAM_DRIVE, INDEX_URL, DOWNLOAD_STATUS_UPDATE_INTERVAL from bot.helper.ext_utils.fs_utils import get_mime_type from bot.helper.ext_utils.bot_utils import * From 7cf7255aabe8e3aa3bff2b66132b47bd0e7433f4 Mon Sep 17 00:00:00 2001 From: jaskaranSM Date: Mon, 13 Jan 2020 19:06:02 +0530 Subject: [PATCH 02/37] Added peers and seeders count for bittorent downloads --- bot/helper/ext_utils/bot_utils.py | 7 ++++++- bot/helper/mirror_utils/download_utils/aria2_download.py | 6 ++++-- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/bot/helper/ext_utils/bot_utils.py b/bot/helper/ext_utils/bot_utils.py index 9f699c7..5f42e6a 100644 --- a/bot/helper/ext_utils/bot_utils.py +++ b/bot/helper/ext_utils/bot_utils.py @@ -90,7 +90,12 @@ def get_readable_message(): if download.status() != MirrorStatus.STATUS_ARCHIVING: msg += f"\n{get_progress_bar_string(download)} {download.progress()} of " \ f"{download.size()}" \ - f" at {download.speed()}, ETA: {download.eta()}\n\n" + f" at {download.speed()}, ETA: {download.eta()} " + if download.status() == MirrorStatus.STATUS_DOWNLOADING: + if hasattr(download,'is_torrent'): + msg += f"| P: {download.download().connections} " \ + f"| S: {download.download().num_seeders}" + msg += "\n\n" return msg diff --git a/bot/helper/mirror_utils/download_utils/aria2_download.py b/bot/helper/mirror_utils/download_utils/aria2_download.py index 974218d..2817ba9 100644 --- a/bot/helper/mirror_utils/download_utils/aria2_download.py +++ b/bot/helper/mirror_utils/download_utils/aria2_download.py @@ -1,4 +1,4 @@ -from bot import aria2 +from bot import aria2,download_dict,download_dict_lock from bot.helper.ext_utils.bot_utils import * from .download_helper import DownloadHelper from bot.helper.mirror_utils.status_utils.aria_download_status import AriaDownloadStatus @@ -26,7 +26,9 @@ def __onDownloadComplete(self, api: API, gid): if self.gid == gid: if api.get_download(gid).followed_by_ids: self.gid = api.get_download(gid).followed_by_ids[0] - download_dict[self._listener.uid] = AriaDownloadStatus(self.gid, self._listener) + with download_dict_lock: + download_dict[self._listener.uid] = AriaDownloadStatus(self.gid, self._listener) + download_dict[self._listener.uid].is_torrent =True update_all_messages() LOGGER.info(f'Changed gid from {gid} to {self.gid}') else: From 457359855d0da57615f14896ff7675df8c77a87c Mon Sep 17 00:00:00 2001 From: lzzy12 Date: Fri, 17 Jan 2020 10:02:46 -0800 Subject: [PATCH 03/37] Tag the author of replied message of mirror command This closes #24 Signed-off-by: lzzy12 --- bot/modules/mirror.py | 24 +++++++++++++----------- 1 file changed, 13 insertions(+), 11 deletions(-) diff --git a/bot/modules/mirror.py b/bot/modules/mirror.py index ad474c9..4243118 100644 --- a/bot/modules/mirror.py +++ b/bot/modules/mirror.py @@ -16,9 +16,10 @@ class MirrorListener(listeners.MirrorListeners): - def __init__(self, bot, update, isTar=False): + def __init__(self, bot, update, isTar=False, tag=None): super().__init__(bot, update) self.isTar = isTar + self.tag = tag def onDownloadStarted(self): pass @@ -98,6 +99,8 @@ def onUploadComplete(self, link: str): if os.path.isdir(f'{DOWNLOAD_DIR}/{self.uid}/{download_dict[self.uid].name()}'): share_url += '/' msg += f'\n\n Shareable link: here' + if self.tag is not None: + msg += f'\ncc: @{self.tag}' try: fs_utils.clean_download(download_dict[self.uid].path()) except FileNotFoundError: @@ -133,19 +136,18 @@ def _mirror(bot, update, isTar=False): link = '' LOGGER.info(link) link = link.strip() - - if len(link) == 0: - if update.message.reply_to_message is not None: - document = update.message.reply_to_message.document - if document is not None and document.mime_type == "application/x-bittorrent": - link = document.get_file().file_path - else: - sendMessage('Only torrent files can be mirrored from telegram', bot, update) - return + reply_to = update.message.reply_to_message + if reply_to is not None: + tag = reply_to.from_user.username + if len(link) == 0: + if reply_to.document is not None and reply_to.document.mime_type == "application/x-bittorrent": + link = reply_to.document.get_file().file_path + else: + tag = None if not bot_utils.is_url(link) and not bot_utils.is_magnet(link): sendMessage('No download source provided', bot, update) return - listener = MirrorListener(bot, update, isTar) + listener = MirrorListener(bot, update, isTar, tag) aria = aria2_download.AriaDownloadHelper(listener) aria.add_download(link, f'{DOWNLOAD_DIR}/{listener.uid}/') sendStatusMessage(update, bot) From 29ee3ac271a59b1eb54b79655a304549ca247d99 Mon Sep 17 00:00:00 2001 From: lzzy12 Date: Sun, 19 Jan 2020 08:50:00 -0800 Subject: [PATCH 04/37] Dockerfile: Move to smaller alpine base image Signed-off-by: lzzy12 --- Dockerfile | 26 +++++++++++++++----------- 1 file changed, 15 insertions(+), 11 deletions(-) diff --git a/Dockerfile b/Dockerfile index 0b26e2d..7cc40ae 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,16 +1,20 @@ -FROM ubuntu:18.04 +FROM python:3-alpine + +WORKDIR /bot +RUN chmod 777 /bot + +# install ca-certificates so that HTTPS works consistently +RUN apk add --no-cache --update ca-certificates aria2 libmagic + +RUN apk add --no-cache --update --virtual .build-deps \ + build-base \ + libffi-dev \ + openssl-dev -WORKDIR /usr/src/app -RUN chmod 777 /usr/src/app -RUN apt -qq update -RUN apt -qq install -y aria2 python3 python3-pip locales COPY requirements.txt . RUN pip3 install --no-cache-dir -r requirements.txt +RUN apk del .build-deps COPY . . -RUN chmod +x aria.sh -RUN locale-gen en_US.UTF-8 -ENV LANG en_US.UTF-8 -ENV LANGUAGE en_US:en -ENV LC_ALL en_US.UTF-8 +RUN chmod +x *.sh -CMD ["bash","start.sh"] +CMD ["sh", "./start.sh"] From 651c893775cd6e52d86c640b7e4bfafbcca78a64 Mon Sep 17 00:00:00 2001 From: lzzy12 Date: Tue, 21 Jan 2020 02:21:38 -0800 Subject: [PATCH 05/37] Minor fix ups --- bot/__init__.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/bot/__init__.py b/bot/__init__.py index 9034147..4063c7f 100644 --- a/bot/__init__.py +++ b/bot/__init__.py @@ -1,11 +1,12 @@ import logging -import aria2p -import threading import os -from dotenv import load_dotenv -import telegram.ext as tg +import threading import time +import aria2p +import telegram.ext as tg +from dotenv import load_dotenv + botStartTime = time.time() if os.path.exists('log.txt'): with open('log.txt', 'r+') as f: @@ -28,7 +29,7 @@ def getConfig(name: str): try: if bool(getConfig('_____REMOVE_THIS_LINE_____')): - logging.ERROR('The README.md file there to be read! Exiting now!') + logging.error('The README.md file there to be read! Exiting now!') exit() except KeyError: pass @@ -80,7 +81,7 @@ def getConfig(name: str): INDEX_URL = None try: IS_TEAM_DRIVE = getConfig('IS_TEAM_DRIVE') - if IS_TEAM_DRIVE == 'True' or IS_TEAM_DRIVE == 'true': + if IS_TEAM_DRIVE.lower() == 'true': IS_TEAM_DRIVE = True else: IS_TEAM_DRIVE = False From 2468b87720b71e67203693fa9f0c07a06188684a Mon Sep 17 00:00:00 2001 From: lzzy12 Date: Tue, 21 Jan 2020 03:49:15 -0800 Subject: [PATCH 06/37] [NOT TESTED] Adding support service accounts --- .gitignore | 1 + bot/__init__.py | 11 +- .../mirror_utils/upload_utils/gdriveTools.py | 155 +++++++++++------- 3 files changed, 103 insertions(+), 64 deletions(-) diff --git a/.gitignore b/.gitignore index b651b32..86facd3 100644 --- a/.gitignore +++ b/.gitignore @@ -10,3 +10,4 @@ data* *.pickle authorized_chats.txt log.txt +accounts/* \ No newline at end of file diff --git a/bot/__init__.py b/bot/__init__.py index 4063c7f..3bc56e3 100644 --- a/bot/__init__.py +++ b/bot/__init__.py @@ -85,9 +85,18 @@ def getConfig(name: str): IS_TEAM_DRIVE = True else: IS_TEAM_DRIVE = False - except KeyError: IS_TEAM_DRIVE = False + +try: + USE_SERVICE_ACCOUNTS = getConfig('USE_SERVICE_ACCOUNTS') + if USE_SERVICE_ACCOUNTS.lower() == 'true': + USE_SERVICE_ACCOUNTS = True + else: + USE_SERVICE_ACCOUNTS = False +except KeyError: + USE_SERVICE_ACCOUNTS = False + updater = tg.Updater(token=BOT_TOKEN) bot = updater.bot dispatcher = updater.dispatcher diff --git a/bot/helper/mirror_utils/upload_utils/gdriveTools.py b/bot/helper/mirror_utils/upload_utils/gdriveTools.py index 7f8f2d7..8eccdad 100644 --- a/bot/helper/mirror_utils/upload_utils/gdriveTools.py +++ b/bot/helper/mirror_utils/upload_utils/gdriveTools.py @@ -1,31 +1,65 @@ -from googleapiclient.discovery import build -from google_auth_oauthlib.flow import InstalledAppFlow +import json +import os +import pickle + from google.auth.transport.requests import Request -from googleapiclient.http import MediaFileUpload +from google_auth_oauthlib.flow import InstalledAppFlow +from googleapiclient.discovery import build from googleapiclient.errors import HttpError -import pickle -import os +from googleapiclient.http import MediaFileUpload +from oauth2client.service_account import ServiceAccountCredentials from tenacity import * -import threading -from bot import LOGGER, parent_id, DOWNLOAD_DIR, IS_TEAM_DRIVE, INDEX_URL, DOWNLOAD_STATUS_UPDATE_INTERVAL -from bot.helper.ext_utils.fs_utils import get_mime_type + +from bot import LOGGER, parent_id, DOWNLOAD_DIR, IS_TEAM_DRIVE, INDEX_URL, DOWNLOAD_STATUS_UPDATE_INTERVAL, \ + USE_SERVICE_ACCOUNTS from bot.helper.ext_utils.bot_utils import * +from bot.helper.ext_utils.fs_utils import get_mime_type logging.getLogger('googleapiclient.discovery').setLevel(logging.ERROR) +G_DRIVE_TOKEN_FILE = "token.pickle" +# Check https://developers.google.com/drive/scopes for all available scopes +OAUTH_SCOPE = ["https://www.googleapis.com/auth/drive"] + +SERVICE_ACCOUNT_INDEX = 0 + + +def authorize(): + # Get credentials + credentials = None + if not USE_SERVICE_ACCOUNTS: + if os.path.exists(G_DRIVE_TOKEN_FILE): + with open(G_DRIVE_TOKEN_FILE, 'rb') as f: + credentials = pickle.load(f) + if credentials is None or not credentials.valid: + if credentials and credentials.expired and credentials.refresh_token: + credentials.refresh(Request()) + else: + flow = InstalledAppFlow.from_client_secrets_file( + 'credentials.json', OAUTH_SCOPE) + LOGGER.info(flow) + credentials = flow.run_console(port=0) + + # Save the credentials for the next run + with open(G_DRIVE_TOKEN_FILE, 'wb') as token: + pickle.dump(credentials, token) + else: + credentials = ServiceAccountCredentials \ + .from_json_keyfile_name(f'accounts/{SERVICE_ACCOUNT_INDEX}.json') + return build('drive', 'v3', credentials=credentials, cache_discovery=False) + + +service = authorize() + class GoogleDriveHelper: + # Redirect URI for installed apps, can be left as is + REDIRECT_URI = "urn:ietf:wg:oauth:2.0:oob" + G_DRIVE_DIR_MIME_TYPE = "application/vnd.google-apps.folder" + G_DRIVE_BASE_DOWNLOAD_URL = "https://drive.google.com/uc?id={}&export=download" def __init__(self, name=None, listener=None): - self.__G_DRIVE_TOKEN_FILE = "token.pickle" - # Check https://developers.google.com/drive/scopes for all available scopes - self.__OAUTH_SCOPE = ["https://www.googleapis.com/auth/drive"] - # Redirect URI for installed apps, can be left as is - self.__REDIRECT_URI = "urn:ietf:wg:oauth:2.0:oob" - self.__G_DRIVE_DIR_MIME_TYPE = "application/vnd.google-apps.folder" - self.__G_DRIVE_BASE_DOWNLOAD_URL = "https://drive.google.com/uc?id={}&export=download" self.__listener = listener - self.__service = self.authorize() self._file_uploaded_bytes = 0 self.uploaded_bytes = 0 self.start_time = 0 @@ -51,7 +85,8 @@ def speed(self): except ZeroDivisionError: return 0 - @retry(wait=wait_exponential(multiplier=2, min=3, max=6),stop=stop_after_attempt(5),retry=retry_if_exception_type(HttpError),before=before_log(LOGGER,logging.DEBUG)) + @retry(wait=wait_exponential(multiplier=2, min=3, max=6), stop=stop_after_attempt(5), + retry=retry_if_exception_type(HttpError), before=before_log(LOGGER, logging.DEBUG)) def _on_upload_progress(self): if self.status is not None: chunk_size = self.status.total_size * self.status.progress() - self._file_uploaded_bytes @@ -60,7 +95,8 @@ def _on_upload_progress(self): self.uploaded_bytes += chunk_size self.total_time += DOWNLOAD_STATUS_UPDATE_INTERVAL - def __upload_empty_file(self, path, file_name, mime_type, parent_id=None): + @staticmethod + def __upload_empty_file(path, file_name, mime_type, parent_id=None): media_body = MediaFileUpload(path, mimetype=mime_type, resumable=False) @@ -71,10 +107,11 @@ def __upload_empty_file(self, path, file_name, mime_type, parent_id=None): } if parent_id is not None: file_metadata['parents'] = [parent_id] - return self.__service.files().create(supportsTeamDrives=True, - body=file_metadata, media_body=media_body).execute() + return service.files().create(supportsTeamDrives=True, + body=file_metadata, media_body=media_body).execute() - @retry(wait=wait_exponential(multiplier=2, min=3, max=6),stop=stop_after_attempt(5),retry=retry_if_exception_type(HttpError),before=before_log(LOGGER,logging.DEBUG)) + @retry(wait=wait_exponential(multiplier=2, min=3, max=6), stop=stop_after_attempt(5), + retry=retry_if_exception_type(HttpError), before=before_log(LOGGER, logging.DEBUG)) def __set_permission(self, drive_id): permissions = { 'role': 'reader', @@ -82,10 +119,13 @@ def __set_permission(self, drive_id): 'value': None, 'withLink': True } - return self.__service.permissions().create(supportsTeamDrives=True, fileId=drive_id, body=permissions).execute() + return service.permissions().create(supportsTeamDrives=True, fileId=drive_id, body=permissions).execute() - @retry(wait=wait_exponential(multiplier=2, min=3, max=6),stop=stop_after_attempt(5),retry=retry_if_exception_type(HttpError),before=before_log(LOGGER,logging.DEBUG)) + @retry(wait=wait_exponential(multiplier=2, min=3, max=6), stop=stop_after_attempt(5), + retry=retry_if_exception_type(HttpError), before=before_log(LOGGER, logging.DEBUG)) def upload_file(self, file_path, file_name, mime_type, parent_id): + global SERVICE_ACCOUNT_INDEX + global service # File body description file_metadata = { 'name': file_name, @@ -99,33 +139,41 @@ def upload_file(self, file_path, file_name, mime_type, parent_id): media_body = MediaFileUpload(file_path, mimetype=mime_type, resumable=False) - response = self.__service.files().create(supportsTeamDrives=True, - body=file_metadata, media_body=media_body).execute() + response = service.files().create(supportsTeamDrives=True, + body=file_metadata, media_body=media_body).execute() if not IS_TEAM_DRIVE: 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')) + drive_file = service.files().get(fileId=response['id']).execute() + download_url = self.G_DRIVE_BASE_DOWNLOAD_URL.format(drive_file.get('id')) return download_url media_body = MediaFileUpload(file_path, mimetype=mime_type, resumable=True, - chunksize=50*1024*1024) + chunksize=50 * 1024 * 1024) # Insert a file - drive_file = self.__service.files().create(supportsTeamDrives=True, - body=file_metadata, media_body=media_body) + drive_file = service.files().create(supportsTeamDrives=True, + body=file_metadata, media_body=media_body) response = None while response is None: if self.is_cancelled: return None - self.status, response = drive_file.next_chunk() + try: + self.status, response = drive_file.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': + SERVICE_ACCOUNT_INDEX += 1 + service = authorize() + raise err self._file_uploaded_bytes = 0 # Insert new permissions if not IS_TEAM_DRIVE: self.__set_permission(response['id']) # Define file instance and get url for download - 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')) + drive_file = 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 def upload(self, file_name: str): @@ -169,15 +217,16 @@ def upload(self, file_name: str): LOGGER.info("Deleting downloaded file/folder..") return link - @retry(wait=wait_exponential(multiplier=2, min=3, max=6),stop=stop_after_attempt(5),retry=retry_if_exception_type(HttpError),before=before_log(LOGGER,logging.DEBUG)) + @retry(wait=wait_exponential(multiplier=2, min=3, max=6), stop=stop_after_attempt(5), + retry=retry_if_exception_type(HttpError), before=before_log(LOGGER, logging.DEBUG)) def create_directory(self, directory_name, parent_id): file_metadata = { "name": directory_name, - "mimeType": self.__G_DRIVE_DIR_MIME_TYPE + "mimeType": self.G_DRIVE_DIR_MIME_TYPE } if parent_id is not None: file_metadata["parents"] = [parent_id] - file = self.__service.files().create(supportsTeamDrives=True, body=file_metadata).execute() + file = service.files().create(supportsTeamDrives=True, body=file_metadata).execute() file_id = file.get("id") if not IS_TEAM_DRIVE: self.__set_permission(file_id) @@ -204,26 +253,6 @@ def upload_dir(self, input_directory, parent_id): new_id = parent_id return new_id - def authorize(self): - # Get credentials - credentials = None - if os.path.exists(self.__G_DRIVE_TOKEN_FILE): - with open(self.__G_DRIVE_TOKEN_FILE, 'rb') as f: - credentials = pickle.load(f) - if credentials is None or not credentials.valid: - if credentials and credentials.expired and credentials.refresh_token: - credentials.refresh(Request()) - else: - flow = InstalledAppFlow.from_client_secrets_file( - 'credentials.json', self.__OAUTH_SCOPE) - LOGGER.info(flow) - credentials = flow.run_console(port=0) - - # Save the credentials for the next run - with open(self.__G_DRIVE_TOKEN_FILE, 'wb') as token: - pickle.dump(credentials, token) - return build('drive', 'v3', credentials=credentials, cache_discovery=False) - def drive_list(self, fileName): msg = "" # Create Search Query for API request. @@ -231,13 +260,13 @@ def drive_list(self, fileName): page_token = None results = [] while True: - response = self.__service.files().list(supportsTeamDrives=True, - includeTeamDriveItems=True, - q=query, - spaces='drive', - fields='nextPageToken, files(id, name, mimeType, size)', - pageToken=page_token, - orderBy='modifiedTime desc').execute() + response = service.files().list(supportsTeamDrives=True, + includeTeamDriveItems=True, + q=query, + spaces='drive', + fields='nextPageToken, files(id, name, mimeType, size)', + pageToken=page_token, + orderBy='modifiedTime desc').execute() for file in response.get('files', []): if len(results) >= 20: break From 134334c15c80f0f0cd3c6a6d4829babf3abf7279 Mon Sep 17 00:00:00 2001 From: lzzy12 Date: Tue, 21 Jan 2020 04:23:23 -0800 Subject: [PATCH 07/37] Added scripts and docs for generating service accounts --- README.md | 31 ++++ add_to_google_group.py | 84 ++++++++++ config_sample.env | 1 + gen_sa_accounts.py | 365 +++++++++++++++++++++++++++++++++++++++++ 4 files changed, 481 insertions(+) create mode 100644 add_to_google_group.py create mode 100644 gen_sa_accounts.py diff --git a/README.md b/README.md index d92f9c4..5e1681e 100644 --- a/README.md +++ b/README.md @@ -84,3 +84,34 @@ sudo docker build . -t mirror-bot ``` sudo docker run mirror-bot ``` + +## Using service accounts for uploading to avoid user rate limit + +Many thanks to [AutoRClone](https://github.com/xyou365/AutoRclone) for the scripts +### Generating service accounts +Step 1. Generate service accounts [What is service account](https://cloud.google.com/iam/docs/service-accounts) [How to use service account in rclone](https://rclone.org/drive/#service-account-support). +--------------------------------- +Let us create only the service accounts that we need. +**Warning:** abuse of this feature is not the aim of autorclone and we do **NOT** recommend that you make a lot of projects, just one project and 100 sa allow you plenty of use, its also possible that overabuse might get your projects banned by google. + +``` +Note: 1 service account can copy around 750gb a day, 1 project makes 100 service accounts so thats 75tb a day, for most users this should easily suffice. +``` + +`python3 gen_sa_accounts.py --quick-setup 1 --new-only` + +A folder named accounts will be created which will contain keys for the service accounts created +``` +We highly recommend to zip this folder and store it somewhere safe, so that you do not have to create a new project everytime you want to deploy the bot +``` +### Adding service accounts to Google Groups: +We use Google Groups to manager our service accounts considering the +[Official limits to the members of Team Drive](https://support.google.com/a/answer/7338880?hl=en) (Limit for individuals and groups directly added as members: 600). + +1. Turn on the Directory API following [official steps](https://developers.google.com/admin-sdk/directory/v1/quickstart/python) (save the generated json file to folder `credentials`). + +2. Create group for your organization [in the Admin console](https://support.google.com/a/answer/33343?hl=en). After create a group, you will have an address for example`sa@yourdomain.com`. + +3. Run `python3 add_to_google_group.py -g sa@yourdomain.com` + +4. Now, add Google Groups (**Step 2**) to manager your service accounts, add the group address `sa@yourdomain.com` or `sa@googlegroups.com` to the Team drive or folder diff --git a/add_to_google_group.py b/add_to_google_group.py new file mode 100644 index 0000000..aaed28b --- /dev/null +++ b/add_to_google_group.py @@ -0,0 +1,84 @@ +# auto rclone +# Add service accounts to groups for your organization +# +# Author Telegram https://t.me/CodyDoby +# Inbox codyd@qq.com + +from __future__ import print_function + +import os +import pickle + +import argparse +import glob +import googleapiclient.discovery +import json +import progress.bar +import time +from google.auth.transport.requests import Request +from google_auth_oauthlib.flow import InstalledAppFlow + +stt = time.time() + +parse = argparse.ArgumentParser( + description='A tool to add service accounts to groups for your organization from a folder containing credential ' + 'files.') +parse.add_argument('--path', '-p', default='accounts', + help='Specify an alternative path to the service accounts folder.') +parse.add_argument('--credentials', '-c', default='credentials/credentials.json', + help='Specify the relative path for the controller file.') +parsereq = parse.add_argument_group('required arguments') +# service-account@googlegroups.com +parsereq.add_argument('--groupaddr', '-g', help='The address of groups for your organization.', required=True) + +args = parse.parse_args() +acc_dir = args.path +gaddr = args.groupaddr +credentials = glob.glob(args.credentials) + +creds = None +if os.path.exists('credentials/token.pickle'): + with open('credentials/token.pickle', 'rb') as token: + creds = pickle.load(token) +# If there are no (valid) credentials available, let the user log in. +if not creds or not creds.valid: + if creds and creds.expired and creds.refresh_token: + creds.refresh(Request()) + else: + flow = InstalledAppFlow.from_client_secrets_file(credentials[0], scopes=[ + 'https://www.googleapis.com/auth/admin.directory.group', + 'https://www.googleapis.com/auth/admin.directory.group.member' + ]) + # creds = flow.run_local_server(port=0) + creds = flow.run_console() + # Save the credentials for the next run + with open('credentials/token.pickle', 'wb') as token: + pickle.dump(creds, token) + +group = googleapiclient.discovery.build("admin", "directory_v1", credentials=creds) + +print(group.members()) + +batch = group.new_batch_http_request() + +sa = glob.glob('%s/*.json' % acc_dir) + +# sa = sa[0:5] + +pbar = progress.bar.Bar("Readying accounts", max=len(sa)) +for i in sa: + ce = json.loads(open(i, 'r').read())['client_email'] + + body = {"email": ce, "role": "MEMBER"} + batch.add(group.members().insert(groupKey=gaddr, body=body)) + # group.members().insert(groupKey=gaddr, body=body).execute() + + pbar.next() +pbar.finish() +print('Adding...') +batch.execute() + +print('Complete.') +hours, rem = divmod((time.time() - stt), 3600) +minutes, sec = divmod(rem, 60) +print("Elapsed Time:\n{:0>2}:{:0>2}:{:05.2f}".format(int(hours), int(minutes), sec)) diff --git a/config_sample.env b/config_sample.env index 2e8252e..686d440 100644 --- a/config_sample.env +++ b/config_sample.env @@ -10,3 +10,4 @@ DOWNLOAD_STATUS_UPDATE_INTERVAL = 5 AUTO_DELETE_MESSAGE_DURATION = 20 IS_TEAM_DRIVE = "" INDEX_URL = "" +USE_SERVICE_ACCOUNTS = "" \ No newline at end of file diff --git a/gen_sa_accounts.py b/gen_sa_accounts.py new file mode 100644 index 0000000..3a11fa9 --- /dev/null +++ b/gen_sa_accounts.py @@ -0,0 +1,365 @@ +from __future__ import print_function + +import errno +import os +import pickle +import sys +from argparse import ArgumentParser +from base64 import b64decode +from glob import glob +from json import loads +from random import choice +from time import sleep + +from google.auth.transport.requests import Request +from google_auth_oauthlib.flow import InstalledAppFlow +from googleapiclient.discovery import build +from googleapiclient.errors import HttpError + +SCOPES = ['https://www.googleapis.com/auth/drive', 'https://www.googleapis.com/auth/cloud-platform', + 'https://www.googleapis.com/auth/iam'] +project_create_ops = [] +current_key_dump = [] +sleep_time = 30 + + +# Create count SAs in project +def _create_accounts(service, project, count): + batch = service.new_batch_http_request(callback=_def_batch_resp) + for i in range(count): + aid = _generate_id('mfc-') + batch.add(service.projects().serviceAccounts().create(name='projects/' + project, body={'accountId': aid, + 'serviceAccount': { + 'displayName': aid}})) + batch.execute() + + +# Create accounts needed to fill project +def _create_remaining_accounts(iam, project): + print('Creating accounts in %s' % project) + sa_count = len(_list_sas(iam, project)) + while sa_count != 100: + _create_accounts(iam, project, 100 - sa_count) + sa_count = len(_list_sas(iam, project)) + + +# Generate a random id +def _generate_id(prefix='saf-'): + chars = '-abcdefghijklmnopqrstuvwxyz1234567890' + return prefix + ''.join(choice(chars) for _ in range(25)) + choice(chars[1:]) + + +# List projects using service +def _get_projects(service): + return [i['projectId'] for i in service.projects().list().execute()['projects']] + + +# Default batch callback handler +def _def_batch_resp(id, resp, exception): + if exception is not None: + if str(exception).startswith(' 0: + current_count = len(_get_projects(cloud)) + if current_count + create_projects <= max_projects: + print('Creating %d projects' % (create_projects)) + nprjs = _create_projects(cloud, create_projects) + selected_projects = nprjs + else: + sys.exit('No, you cannot create %d new project (s).\n' + 'Please reduce value of --quick-setup.\n' + 'Remember that you can totally create %d projects (%d already).\n' + 'Please do not delete existing projects unless you know what you are doing' % ( + create_projects, max_projects, current_count)) + else: + print('Will overwrite all service accounts in existing projects.\n' + 'So make sure you have some projects already.') + input("Press Enter to continue...") + + if enable_services: + ste = [] + ste.append(enable_services) + if enable_services == '~': + ste = selected_projects + elif enable_services == '*': + ste = _get_projects(cloud) + services = [i + '.googleapis.com' for i in services] + print('Enabling services') + _enable_services(serviceusage, ste, services) + if create_sas: + stc = [] + stc.append(create_sas) + if create_sas == '~': + stc = selected_projects + elif create_sas == '*': + stc = _get_projects(cloud) + for i in stc: + _create_remaining_accounts(iam, i) + if download_keys: + try: + os.mkdir(path) + except OSError as e: + if e.errno == errno.EEXIST: + pass + else: + raise + std = [] + std.append(download_keys) + if download_keys == '~': + std = selected_projects + elif download_keys == '*': + std = _get_projects(cloud) + _create_sa_keys(iam, std, path) + if delete_sas: + std = [] + std.append(delete_sas) + if delete_sas == '~': + std = selected_projects + elif delete_sas == '*': + std = _get_projects(cloud) + for i in std: + print('Deleting service accounts in %s' % i) + _delete_sas(iam, i) + + +if __name__ == '__main__': + parse = ArgumentParser(description='A tool to create Google service accounts.') + parse.add_argument('--path', '-p', default='accounts', + help='Specify an alternate directory to output the credential files.') + parse.add_argument('--token', default='token.pickle', help='Specify the pickle token file path.') + parse.add_argument('--credentials', default='credentials.json', help='Specify the credentials file path.') + parse.add_argument('--list-projects', default=False, action='store_true', + help='List projects viewable by the user.') + parse.add_argument('--list-sas', default=False, help='List service accounts in a project.') + parse.add_argument('--create-projects', type=int, default=None, help='Creates up to N projects.') + parse.add_argument('--max-projects', type=int, default=12, help='Max amount of project allowed. Default: 12') + parse.add_argument('--enable-services', default=None, + help='Enables services on the project. Default: IAM and Drive') + parse.add_argument('--services', nargs='+', default=['iam', 'drive'], + help='Specify a different set of services to enable. Overrides the default.') + parse.add_argument('--create-sas', default=None, help='Create service accounts in a project.') + parse.add_argument('--delete-sas', default=None, help='Delete service accounts in a project.') + parse.add_argument('--download-keys', default=None, help='Download keys for all the service accounts in a project.') + parse.add_argument('--quick-setup', default=None, type=int, + help='Create projects, enable services, create service accounts and download keys. ') + parse.add_argument('--new-only', default=False, action='store_true', help='Do not use exisiting projects.') + args = parse.parse_args() + # If credentials file is invalid, search for one. + if not os.path.exists(args.credentials): + options = glob('*.json') + print('No credentials found at %s. Please enable the Drive API in:\n' + 'https://developers.google.com/drive/api/v3/quickstart/python\n' + 'and save the json file as credentials.json' % args.credentials) + if len(options) < 1: + exit(-1) + else: + i = 0 + print('Select a credentials file below.') + inp_options = [str(i) for i in list(range(1, len(options) + 1))] + options + while i < len(options): + print(' %d) %s' % (i + 1, options[i])) + i += 1 + inp = None + while True: + inp = input('> ') + if inp in inp_options: + break + if inp in options: + args.credentials = inp + else: + args.credentials = options[int(inp) - 1] + print('Use --credentials %s next time to use this credentials file.' % args.credentials) + if args.quick_setup: + opt = '*' + if args.new_only: + opt = '~' + args.services = ['iam', 'drive'] + args.create_projects = args.quick_setup + args.enable_services = opt + args.create_sas = opt + args.download_keys = opt + resp = serviceaccountfactory( + path=args.path, + token=args.token, + credentials=args.credentials, + list_projects=args.list_projects, + list_sas=args.list_sas, + create_projects=args.create_projects, + max_projects=args.max_projects, + create_sas=args.create_sas, + delete_sas=args.delete_sas, + enable_services=args.enable_services, + services=args.services, + download_keys=args.download_keys + ) + if resp is not None: + if args.list_projects: + if resp: + print('Projects (%d):' % len(resp)) + for i in resp: + print(' ' + i) + else: + print('No projects.') + elif args.list_sas: + if resp: + print('Service accounts in %s (%d):' % (args.list_sas, len(resp))) + for i in resp: + print(' %s (%s)' % (i['email'], i['uniqueId'])) + else: + print('No service accounts.') From be68169065cbeb586606bb91b0c563f5e717e8d4 Mon Sep 17 00:00:00 2001 From: lzzy12 Date: Tue, 21 Jan 2020 04:30:26 -0800 Subject: [PATCH 08/37] gen_sa_accounts: Save credentials with indexed file name --- gen_sa_accounts.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/gen_sa_accounts.py b/gen_sa_accounts.py index 3a11fa9..c77ef65 100644 --- a/gen_sa_accounts.py +++ b/gen_sa_accounts.py @@ -1,5 +1,3 @@ -from __future__ import print_function - import errno import os import pickle @@ -147,9 +145,11 @@ def _create_sa_keys(iam, projects, path): print('Redownloading keys from %s' % i) current_key_dump = [] else: + index = 0 for j in current_key_dump: - with open('%s/%s.json' % (path, j[0]), 'w+') as f: + with open(f'{path}/{index}.json', 'w+') as f: f.write(j[1]) + index += 1 # Delete Service Accounts From f180b90e14c27e72bccf8a9b494bab6060036098 Mon Sep 17 00:00:00 2001 From: lzzy12 Date: Thu, 23 Jan 2020 00:03:59 -0800 Subject: [PATCH 09/37] gdriveTools: Avoid using oauth2 library for service accounts oauth2 library is deprecated --- bot/helper/mirror_utils/upload_utils/gdriveTools.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/bot/helper/mirror_utils/upload_utils/gdriveTools.py b/bot/helper/mirror_utils/upload_utils/gdriveTools.py index 8eccdad..61feb53 100644 --- a/bot/helper/mirror_utils/upload_utils/gdriveTools.py +++ b/bot/helper/mirror_utils/upload_utils/gdriveTools.py @@ -3,11 +3,11 @@ import pickle from google.auth.transport.requests import Request +from google.oauth2 import service_account from google_auth_oauthlib.flow import InstalledAppFlow from googleapiclient.discovery import build from googleapiclient.errors import HttpError from googleapiclient.http import MediaFileUpload -from oauth2client.service_account import ServiceAccountCredentials from tenacity import * from bot import LOGGER, parent_id, DOWNLOAD_DIR, IS_TEAM_DRIVE, INDEX_URL, DOWNLOAD_STATUS_UPDATE_INTERVAL, \ @@ -44,8 +44,9 @@ def authorize(): with open(G_DRIVE_TOKEN_FILE, 'wb') as token: pickle.dump(credentials, token) else: - credentials = ServiceAccountCredentials \ - .from_json_keyfile_name(f'accounts/{SERVICE_ACCOUNT_INDEX}.json') + credentials = service_account.Credentials \ + .from_service_account_file(f'accounts/{SERVICE_ACCOUNT_INDEX}.json', + scopes=OAUTH_SCOPE) return build('drive', 'v3', credentials=credentials, cache_discovery=False) From ba4b1d9c482a43f7494852bbc73298b598473a26 Mon Sep 17 00:00:00 2001 From: lzzy12 Date: Thu, 23 Jan 2020 00:13:51 -0800 Subject: [PATCH 10/37] gdriveTools: Fixed update Interval oauth2 library is deprecated --- bot/helper/mirror_utils/upload_utils/gdriveTools.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/bot/helper/mirror_utils/upload_utils/gdriveTools.py b/bot/helper/mirror_utils/upload_utils/gdriveTools.py index 61feb53..5bc1500 100644 --- a/bot/helper/mirror_utils/upload_utils/gdriveTools.py +++ b/bot/helper/mirror_utils/upload_utils/gdriveTools.py @@ -10,8 +10,7 @@ from googleapiclient.http import MediaFileUpload from tenacity import * -from bot import LOGGER, parent_id, DOWNLOAD_DIR, IS_TEAM_DRIVE, INDEX_URL, DOWNLOAD_STATUS_UPDATE_INTERVAL, \ - USE_SERVICE_ACCOUNTS +from bot import LOGGER, parent_id, DOWNLOAD_DIR, IS_TEAM_DRIVE, INDEX_URL, USE_SERVICE_ACCOUNTS from bot.helper.ext_utils.bot_utils import * from bot.helper.ext_utils.fs_utils import get_mime_type @@ -58,6 +57,7 @@ class GoogleDriveHelper: REDIRECT_URI = "urn:ietf:wg:oauth:2.0:oob" G_DRIVE_DIR_MIME_TYPE = "application/vnd.google-apps.folder" G_DRIVE_BASE_DOWNLOAD_URL = "https://drive.google.com/uc?id={}&export=download" + UPDATE_INTERVAL = 5 def __init__(self, name=None, listener=None): self.__listener = listener @@ -94,7 +94,7 @@ def _on_upload_progress(self): self._file_uploaded_bytes = self.status.total_size * self.status.progress() LOGGER.info(f'Chunk size: {get_readable_file_size(chunk_size)}') self.uploaded_bytes += chunk_size - self.total_time += DOWNLOAD_STATUS_UPDATE_INTERVAL + self.total_time += self.UPDATE_INTERVAL @staticmethod def __upload_empty_file(path, file_name, mime_type, parent_id=None): @@ -183,7 +183,7 @@ def upload(self, file_name: str): file_path = f"{file_dir}/{file_name}" LOGGER.info("Uploading File: " + file_path) self.start_time = time.time() - self.updater = setInterval(5, self._on_upload_progress) + self.updater = setInterval(self.UPDATE_INTERVAL, self._on_upload_progress) if os.path.isfile(file_path): try: mime_type = get_mime_type(file_path) From ddb44ee3038586c2f43692b70d8d94155dfcb7ca Mon Sep 17 00:00:00 2001 From: lzzy12 Date: Thu, 23 Jan 2020 01:11:00 -0800 Subject: [PATCH 11/37] Fixed crash if no url or magnet is provided to mirror This fixes #32 Signed-off-by: lzzy12 --- .../download_utils/aria2_download.py | 22 ++++++++++++------- 1 file changed, 14 insertions(+), 8 deletions(-) diff --git a/bot/helper/mirror_utils/download_utils/aria2_download.py b/bot/helper/mirror_utils/download_utils/aria2_download.py index 2817ba9..e31b6aa 100644 --- a/bot/helper/mirror_utils/download_utils/aria2_download.py +++ b/bot/helper/mirror_utils/download_utils/aria2_download.py @@ -1,10 +1,12 @@ -from bot import aria2,download_dict,download_dict_lock +from aria2p import API +from aria2p.client import ClientException + +from bot import aria2 from bot.helper.ext_utils.bot_utils import * -from .download_helper import DownloadHelper from bot.helper.mirror_utils.status_utils.aria_download_status import AriaDownloadStatus from bot.helper.telegram_helper.message_utils import * -import threading -from aria2p import API +from .download_helper import DownloadHelper + class AriaDownloadHelper(DownloadHelper): @@ -53,10 +55,14 @@ def __onDownloadError(self, api, gid): self._listener.onDownloadError(error) def add_download(self, link: str, path): - if is_magnet(link): - download = aria2.add_magnet(link, {'dir': path}) - else: - download = aria2.add_uris([link], {'dir': path}) + try: + if is_magnet(link): + download = aria2.add_magnet(link, {'dir': path}) + else: + download = aria2.add_uris([link], {'dir': path}) + except ClientException as err: + self._listener.onDownloadError(err.message) + return self.gid = download.gid with download_dict_lock: download_dict[self._listener.uid] = AriaDownloadStatus(self.gid, self._listener) From ddf809a85fd6d4df9ee76a0f4a7da5c1cb75363e Mon Sep 17 00:00:00 2001 From: lzzy12 Date: Thu, 23 Jan 2020 01:30:25 -0800 Subject: [PATCH 12/37] Fixed wrong error message being sent on download timeout Signed-off-by: lzzy12 --- bot/helper/mirror_utils/download_utils/aria2_download.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/bot/helper/mirror_utils/download_utils/aria2_download.py b/bot/helper/mirror_utils/download_utils/aria2_download.py index e31b6aa..5063fd7 100644 --- a/bot/helper/mirror_utils/download_utils/aria2_download.py +++ b/bot/helper/mirror_utils/download_utils/aria2_download.py @@ -39,12 +39,16 @@ def __onDownloadComplete(self, api: API, gid): def __onDownloadPause(self, api, gid): if self.gid == gid: LOGGER.info("Called onDownloadPause") - self._listener.onDownloadError('Download stopped by user!') + download = api.get_download(gid) + error = download.error_message + self._listener.onDownloadError(error) def __onDownloadStopped(self, api, gid): if self.gid == gid: LOGGER.info("Called on_download_stop") - self._listener.onDownloadError('Download stopped by user!') + download = api.get_download(gid) + error = download.error_message + self._listener.onDownloadError(error) def __onDownloadError(self, api, gid): with self._resource_lock: From e4d5936ca5f7c6261a7935e2a94fc96823a537d3 Mon Sep 17 00:00:00 2001 From: jaskaranSM Date: Sun, 2 Feb 2020 14:48:23 +0530 Subject: [PATCH 13/37] Rewrite tar function - shutil changes pwd while archiving. as some functions of bot require relative paths, it is neccesary to keep pwd as bot root dir. --- bot/helper/ext_utils/fs_utils.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/bot/helper/ext_utils/fs_utils.py b/bot/helper/ext_utils/fs_utils.py index ad3b40f..36955f1 100644 --- a/bot/helper/ext_utils/fs_utils.py +++ b/bot/helper/ext_utils/fs_utils.py @@ -4,6 +4,7 @@ import os import pathlib import mimetypes +import tarfile def clean_download(path: str): @@ -30,12 +31,14 @@ def exit_clean_up(signal, frame): sys.exit(1) -def tar(orig_path: str): - path = pathlib.PurePath(orig_path) - base = path.name - root = pathlib.Path(path.parent.as_posix()).absolute().as_posix() - LOGGER.info(f'Tar: orig_path: {orig_path}, base: {base}, root: {root}') - return shutil.make_archive(orig_path, 'tar', root, base) +def tar(org_path): + tar_path = org_path + ".tar" + path = pathlib.PurePath(org_path) + LOGGER.info(f'Tar: orig_path: {org_path}, tar_path: {tar_path}') + tar = tarfile.open(tar_path, "w") + tar.add(org_path,arcname=path.name) + tar.close() + return tar_path def get_mime_type(file_path): From c6566c767dc2427f8212afbdb7268722d4ef29f6 Mon Sep 17 00:00:00 2001 From: jaskaranSM Date: Mon, 3 Feb 2020 17:05:46 +0530 Subject: [PATCH 14/37] Refactor service account implementation --- .../mirror_utils/upload_utils/gdriveTools.py | 132 +++++++++--------- 1 file changed, 66 insertions(+), 66 deletions(-) diff --git a/bot/helper/mirror_utils/upload_utils/gdriveTools.py b/bot/helper/mirror_utils/upload_utils/gdriveTools.py index 5bc1500..fbe3d92 100644 --- a/bot/helper/mirror_utils/upload_utils/gdriveTools.py +++ b/bot/helper/mirror_utils/upload_utils/gdriveTools.py @@ -15,51 +15,19 @@ from bot.helper.ext_utils.fs_utils import get_mime_type logging.getLogger('googleapiclient.discovery').setLevel(logging.ERROR) - -G_DRIVE_TOKEN_FILE = "token.pickle" -# Check https://developers.google.com/drive/scopes for all available scopes -OAUTH_SCOPE = ["https://www.googleapis.com/auth/drive"] - SERVICE_ACCOUNT_INDEX = 0 - -def authorize(): - # Get credentials - credentials = None - if not USE_SERVICE_ACCOUNTS: - if os.path.exists(G_DRIVE_TOKEN_FILE): - with open(G_DRIVE_TOKEN_FILE, 'rb') as f: - credentials = pickle.load(f) - if credentials is None or not credentials.valid: - if credentials and credentials.expired and credentials.refresh_token: - credentials.refresh(Request()) - else: - flow = InstalledAppFlow.from_client_secrets_file( - 'credentials.json', OAUTH_SCOPE) - LOGGER.info(flow) - credentials = flow.run_console(port=0) - - # Save the credentials for the next run - with open(G_DRIVE_TOKEN_FILE, 'wb') as token: - pickle.dump(credentials, token) - else: - credentials = service_account.Credentials \ - .from_service_account_file(f'accounts/{SERVICE_ACCOUNT_INDEX}.json', - scopes=OAUTH_SCOPE) - return build('drive', 'v3', credentials=credentials, cache_discovery=False) - - -service = authorize() - - class GoogleDriveHelper: - # Redirect URI for installed apps, can be left as is - REDIRECT_URI = "urn:ietf:wg:oauth:2.0:oob" - G_DRIVE_DIR_MIME_TYPE = "application/vnd.google-apps.folder" - G_DRIVE_BASE_DOWNLOAD_URL = "https://drive.google.com/uc?id={}&export=download" - UPDATE_INTERVAL = 5 - def __init__(self, name=None, listener=None): + self.__G_DRIVE_TOKEN_FILE = "token.pickle" + # Check https://developers.google.com/drive/scopes for all available scopes + self.__OAUTH_SCOPE = ['https://www.googleapis.com/auth/drive'] + # Redirect URI for installed apps, can be left as is + self.__REDIRECT_URI = "urn:ietf:wg:oauth:2.0:oob" + self.__G_DRIVE_DIR_MIME_TYPE = "application/vnd.google-apps.folder" + self.__G_DRIVE_BASE_DOWNLOAD_URL = "https://drive.google.com/uc?id={}&export=download" + self.__listener = listener + self.__service = self.authorize() self.__listener = listener self._file_uploaded_bytes = 0 self.uploaded_bytes = 0 @@ -108,7 +76,7 @@ def __upload_empty_file(path, file_name, mime_type, parent_id=None): } if parent_id is not None: file_metadata['parents'] = [parent_id] - return service.files().create(supportsTeamDrives=True, + return self.__service.files().create(supportsTeamDrives=True, body=file_metadata, media_body=media_body).execute() @retry(wait=wait_exponential(multiplier=2, min=3, max=6), stop=stop_after_attempt(5), @@ -120,13 +88,11 @@ def __set_permission(self, drive_id): 'value': None, 'withLink': True } - return service.permissions().create(supportsTeamDrives=True, fileId=drive_id, body=permissions).execute() + return self.__service.permissions().create(supportsTeamDrives=True, fileId=drive_id, body=permissions).execute() @retry(wait=wait_exponential(multiplier=2, min=3, max=6), stop=stop_after_attempt(5), retry=retry_if_exception_type(HttpError), before=before_log(LOGGER, logging.DEBUG)) def upload_file(self, file_path, file_name, mime_type, parent_id): - global SERVICE_ACCOUNT_INDEX - global service # File body description file_metadata = { 'name': file_name, @@ -140,12 +106,12 @@ def upload_file(self, file_path, file_name, mime_type, parent_id): media_body = MediaFileUpload(file_path, mimetype=mime_type, resumable=False) - response = service.files().create(supportsTeamDrives=True, + response = self.__service.files().create(supportsTeamDrives=True, body=file_metadata, media_body=media_body).execute() if not IS_TEAM_DRIVE: self.__set_permission(response['id']) - drive_file = service.files().get(fileId=response['id']).execute() - download_url = self.G_DRIVE_BASE_DOWNLOAD_URL.format(drive_file.get('id')) + drive_file = self.__service.files().get(fileId=response['id']).execute() + download_url = self.__G_DRIVE_BASE_DOWNLOAD_URL.format(drive_file.get('id')) return download_url media_body = MediaFileUpload(file_path, mimetype=mime_type, @@ -153,7 +119,7 @@ def upload_file(self, file_path, file_name, mime_type, parent_id): chunksize=50 * 1024 * 1024) # Insert a file - drive_file = service.files().create(supportsTeamDrives=True, + drive_file = self.__service.files().create(supportsTeamDrives=True, body=file_metadata, media_body=media_body) response = None while response is None: @@ -165,16 +131,18 @@ def upload_file(self, file_path, file_name, mime_type, parent_id): if err.resp.get('content-type', '').startswith('application/json'): reason = json.loads(err.content).get('error').get('errors')[0].get('reason') if reason == 'userRateLimitExceeded': + global SERVICE_ACCOUNT_INDEX SERVICE_ACCOUNT_INDEX += 1 - service = authorize() + LOGGER.info(f"Switching to {SERVICE_ACCOUNT_INDEX}.json service account") + self.__service = self.authorize() raise err self._file_uploaded_bytes = 0 # Insert new permissions if not IS_TEAM_DRIVE: self.__set_permission(response['id']) # Define file instance and get url for download - drive_file = service.files().get(supportsTeamDrives=True, fileId=response['id']).execute() - download_url = self.G_DRIVE_BASE_DOWNLOAD_URL.format(drive_file.get('id')) + 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 def upload(self, file_name: str): @@ -192,9 +160,13 @@ def upload(self, file_name: str): raise Exception('Upload has been manually cancelled') LOGGER.info("Uploaded To G-Drive: " + file_path) except Exception as e: - LOGGER.info(f"Total Attempts: {e.last_attempt.attempt_number}") - LOGGER.error(e.last_attempt.exception()) - self.__listener.onUploadError(e) + if isinstance(e,RetryError): + LOGGER.info(f"Total Attempts: {e.last_attempt.attempt_number}") + err = e.last_attempt.exception() + else: + err = e + LOGGER.error(err) + self.__listener.onUploadError(str(err)) return finally: self.updater.cancel() @@ -207,9 +179,13 @@ def upload(self, file_name: str): LOGGER.info("Uploaded To G-Drive: " + file_name) link = f"https://drive.google.com/folderview?id={dir_id}" except Exception as e: - LOGGER.info(f"Total Attempts: {e.last_attempt.attempt_number}") - LOGGER.error(e.last_attempt.exception()) - self.__listener.onUploadError(e) + if isinstance(e,RetryError): + LOGGER.info(f"Total Attempts: {e.last_attempt.attempt_number}") + err = e.last_attempt.exception() + else: + err = e + LOGGER.error(err) + self.__listener.onUploadError(str(err)) return finally: self.updater.cancel() @@ -223,11 +199,11 @@ def upload(self, file_name: str): def create_directory(self, directory_name, parent_id): file_metadata = { "name": directory_name, - "mimeType": self.G_DRIVE_DIR_MIME_TYPE + "mimeType": self.__G_DRIVE_DIR_MIME_TYPE } if parent_id is not None: file_metadata["parents"] = [parent_id] - file = service.files().create(supportsTeamDrives=True, body=file_metadata).execute() + file = self.__service.files().create(supportsTeamDrives=True, body=file_metadata).execute() file_id = file.get("id") if not IS_TEAM_DRIVE: self.__set_permission(file_id) @@ -254,14 +230,39 @@ def upload_dir(self, input_directory, parent_id): new_id = parent_id return new_id + def authorize(self): + # Get credentials + credentials = None + if not USE_SERVICE_ACCOUNTS: + if os.path.exists(G_DRIVE_TOKEN_FILE): + with open(G_DRIVE_TOKEN_FILE, 'rb') as f: + credentials = pickle.load(f) + if credentials is None or not credentials.valid: + if credentials and credentials.expired and credentials.refresh_token: + credentials.refresh(Request()) + else: + flow = InstalledAppFlow.from_client_secrets_file( + 'credentials.json', self.__OAUTH_SCOPE) + LOGGER.info(flow) + credentials = flow.run_console(port=0) + + # Save the credentials for the next run + with open(G_DRIVE_TOKEN_FILE, 'wb') as token: + pickle.dump(credentials, token) + else: + LOGGER.info(f"Authorizing with {SERVICE_ACCOUNT_INDEX}.json service account") + credentials = service_account.Credentials.from_service_account_file(f'accounts/{SERVICE_ACCOUNT_INDEX}.json', + scopes=self.__OAUTH_SCOPE) + return build('drive','v3', credentials=credentials, cache_discovery=False) + def drive_list(self, fileName): msg = "" # Create Search Query for API request. query = f"'{parent_id}' in parents and (name contains '{fileName}')" page_token = None - results = [] + count = 0 while True: - response = service.files().list(supportsTeamDrives=True, + response = self.__service.files().list(supportsTeamDrives=True, includeTeamDriveItems=True, q=query, spaces='drive', @@ -269,7 +270,7 @@ def drive_list(self, fileName): pageToken=page_token, orderBy='modifiedTime desc').execute() for file in response.get('files', []): - if len(results) >= 20: + if count >= 20: break if file.get( 'mimeType') == "application/vnd.google-apps.folder": # Detect Whether Current Entity is a Folder or File. @@ -285,9 +286,8 @@ def drive_list(self, fileName): url = f'{INDEX_URL}/{file.get("name")}' msg += f' | Index URL' msg += '\n' - results.append(file) + count += 1 page_token = response.get('nextPageToken', None) - if page_token is None: + if page_token is None or count >= 20: break - del results return msg From a7cfcf2cc17b63c7afa16f432bf94c077a18c3cc Mon Sep 17 00:00:00 2001 From: jaskaranSM <37726998+jaskaranSM@users.noreply.github.com> Date: Tue, 4 Feb 2020 21:23:56 +0530 Subject: [PATCH 15/37] Added missing property --- bot/helper/mirror_utils/upload_utils/gdriveTools.py | 1 + 1 file changed, 1 insertion(+) diff --git a/bot/helper/mirror_utils/upload_utils/gdriveTools.py b/bot/helper/mirror_utils/upload_utils/gdriveTools.py index fbe3d92..453247f 100644 --- a/bot/helper/mirror_utils/upload_utils/gdriveTools.py +++ b/bot/helper/mirror_utils/upload_utils/gdriveTools.py @@ -31,6 +31,7 @@ def __init__(self, name=None, listener=None): self.__listener = listener self._file_uploaded_bytes = 0 self.uploaded_bytes = 0 + self.UPDATE_INTERVAL = 5 self.start_time = 0 self.total_time = 0 self._should_update = True From 97d6fc2921b309eed2b72b2355631eef1a7d4e8c Mon Sep 17 00:00:00 2001 From: Shivam Jha Date: Wed, 5 Feb 2020 01:54:11 -0800 Subject: [PATCH 16/37] UploadStatus: Handle exception ZeroDivisionError --- bot/helper/mirror_utils/status_utils/upload_status.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/bot/helper/mirror_utils/status_utils/upload_status.py b/bot/helper/mirror_utils/status_utils/upload_status.py index a7e7c60..7fad343 100644 --- a/bot/helper/mirror_utils/status_utils/upload_status.py +++ b/bot/helper/mirror_utils/status_utils/upload_status.py @@ -28,8 +28,10 @@ def name(self): return self.obj.name def progress_raw(self): - return self.obj.uploaded_bytes / self.__size * 100 - + try: + return self.obj.uploaded_bytes / self.__size * 100 + except ZeroDivisionError: + return 0 def progress(self): return f'{round(self.progress_raw(), 2)}%' From 8b9a73ab887bfbfe9fa0000ca8c5aee15d6f505d Mon Sep 17 00:00:00 2001 From: lzzy12 Date: Wed, 5 Feb 2020 02:32:27 -0800 Subject: [PATCH 17/37] Fixed upload of empty files for Team Drive Signed-off-by: lzzy12 --- .../mirror_utils/upload_utils/gdriveTools.py | 20 +++---------------- 1 file changed, 3 insertions(+), 17 deletions(-) diff --git a/bot/helper/mirror_utils/upload_utils/gdriveTools.py b/bot/helper/mirror_utils/upload_utils/gdriveTools.py index 453247f..b27d940 100644 --- a/bot/helper/mirror_utils/upload_utils/gdriveTools.py +++ b/bot/helper/mirror_utils/upload_utils/gdriveTools.py @@ -65,21 +65,6 @@ def _on_upload_progress(self): self.uploaded_bytes += chunk_size self.total_time += self.UPDATE_INTERVAL - @staticmethod - def __upload_empty_file(path, file_name, mime_type, parent_id=None): - media_body = MediaFileUpload(path, - mimetype=mime_type, - resumable=False) - file_metadata = { - 'name': file_name, - 'description': 'mirror', - 'mimeType': mime_type, - } - if parent_id is not None: - file_metadata['parents'] = [parent_id] - return self.__service.files().create(supportsTeamDrives=True, - body=file_metadata, media_body=media_body).execute() - @retry(wait=wait_exponential(multiplier=2, min=3, max=6), stop=stop_after_attempt(5), retry=retry_if_exception_type(HttpError), before=before_log(LOGGER, logging.DEBUG)) def __set_permission(self, drive_id): @@ -108,10 +93,11 @@ def upload_file(self, file_path, file_name, mime_type, parent_id): mimetype=mime_type, resumable=False) response = self.__service.files().create(supportsTeamDrives=True, - body=file_metadata, media_body=media_body).execute() + body=file_metadata, media_body=media_body).execute() if not IS_TEAM_DRIVE: self.__set_permission(response['id']) - 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 media_body = MediaFileUpload(file_path, From 5057be2f291ed35064b48201c57393abd1cdb5e6 Mon Sep 17 00:00:00 2001 From: jaskaranSM Date: Sun, 16 Feb 2020 12:29:51 +0530 Subject: [PATCH 18/37] Handle TimedOut Errors --- bot/__init__.py | 2 +- bot/helper/telegram_helper/message_utils.py | 16 +++++++++++----- 2 files changed, 12 insertions(+), 6 deletions(-) diff --git a/bot/__init__.py b/bot/__init__.py index 3bc56e3..8180995 100644 --- a/bot/__init__.py +++ b/bot/__init__.py @@ -97,6 +97,6 @@ def getConfig(name: str): except KeyError: USE_SERVICE_ACCOUNTS = False -updater = tg.Updater(token=BOT_TOKEN) +updater = tg.Updater(token=BOT_TOKEN,request_kwargs={'read_timeout': 60, 'connect_timeout': 60}) bot = updater.bot dispatcher = updater.dispatcher diff --git a/bot/helper/telegram_helper/message_utils.py b/bot/helper/telegram_helper/message_utils.py index ad4e6fe..9a15264 100644 --- a/bot/helper/telegram_helper/message_utils.py +++ b/bot/helper/telegram_helper/message_utils.py @@ -9,25 +9,31 @@ def sendMessage(text: str, bot, update: Update): - return bot.send_message(update.message.chat_id, + try: + return bot.send_message(update.message.chat_id, + reply_to_message_id=update.message.message_id, + text=text, parse_mode='HTMl',timeout=60) + except TimedOut as e: + LOGGER.error(str(e)) + time.sleep(5) #sleep for some time and try request again + return bot.send_message(update.message.chat_id, reply_to_message_id=update.message.message_id, - text=text, parse_mode='HTMl') + text=text, parse_mode='HTMl',timeout=60) def editMessage(text: str, message: Message): try: bot.edit_message_text(text=text, message_id=message.message_id, chat_id=message.chat.id, - parse_mode='HTMl') + parse_mode='HTMl',timeout=60) except TimedOut as e: LOGGER.error(str(e)) - pass def deleteMessage(bot, message: Message): try: bot.delete_message(chat_id=message.chat.id, - message_id=message.message_id) + message_id=message.message_id,timeout=60) except Exception as e: LOGGER.error(str(e)) From df82047f79b89227572bbfc2dcba91504ad6d6c5 Mon Sep 17 00:00:00 2001 From: jaskaranSM Date: Sun, 16 Feb 2020 12:31:10 +0530 Subject: [PATCH 19/37] Not needed --- bot/helper/telegram_helper/message_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/bot/helper/telegram_helper/message_utils.py b/bot/helper/telegram_helper/message_utils.py index 9a15264..6dd9e9b 100644 --- a/bot/helper/telegram_helper/message_utils.py +++ b/bot/helper/telegram_helper/message_utils.py @@ -25,7 +25,7 @@ def editMessage(text: str, message: Message): try: bot.edit_message_text(text=text, message_id=message.message_id, chat_id=message.chat.id, - parse_mode='HTMl',timeout=60) + parse_mode='HTMl') except TimedOut as e: LOGGER.error(str(e)) @@ -33,7 +33,7 @@ def editMessage(text: str, message: Message): def deleteMessage(bot, message: Message): try: bot.delete_message(chat_id=message.chat.id, - message_id=message.message_id,timeout=60) + message_id=message.message_id) except Exception as e: LOGGER.error(str(e)) From bb3a13cf2f32b716632c6ed1a34b9675477cd02b Mon Sep 17 00:00:00 2001 From: jaskaranSM Date: Mon, 17 Feb 2020 21:04:48 +0530 Subject: [PATCH 20/37] Revert "Not needed" This reverts commit df82047f79b89227572bbfc2dcba91504ad6d6c5. --- bot/helper/telegram_helper/message_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/bot/helper/telegram_helper/message_utils.py b/bot/helper/telegram_helper/message_utils.py index 6dd9e9b..9a15264 100644 --- a/bot/helper/telegram_helper/message_utils.py +++ b/bot/helper/telegram_helper/message_utils.py @@ -25,7 +25,7 @@ def editMessage(text: str, message: Message): try: bot.edit_message_text(text=text, message_id=message.message_id, chat_id=message.chat.id, - parse_mode='HTMl') + parse_mode='HTMl',timeout=60) except TimedOut as e: LOGGER.error(str(e)) @@ -33,7 +33,7 @@ def editMessage(text: str, message: Message): def deleteMessage(bot, message: Message): try: bot.delete_message(chat_id=message.chat.id, - message_id=message.message_id) + message_id=message.message_id,timeout=60) except Exception as e: LOGGER.error(str(e)) From 869711123bacd7f908d288a242cbdf2bd251c573 Mon Sep 17 00:00:00 2001 From: jaskaranSM Date: Mon, 17 Feb 2020 21:05:35 +0530 Subject: [PATCH 21/37] Revert "Handle TimedOut Errors" This reverts commit 5057be2f291ed35064b48201c57393abd1cdb5e6. --- bot/__init__.py | 2 +- bot/helper/telegram_helper/message_utils.py | 16 +++++----------- 2 files changed, 6 insertions(+), 12 deletions(-) diff --git a/bot/__init__.py b/bot/__init__.py index 8180995..3bc56e3 100644 --- a/bot/__init__.py +++ b/bot/__init__.py @@ -97,6 +97,6 @@ def getConfig(name: str): except KeyError: USE_SERVICE_ACCOUNTS = False -updater = tg.Updater(token=BOT_TOKEN,request_kwargs={'read_timeout': 60, 'connect_timeout': 60}) +updater = tg.Updater(token=BOT_TOKEN) bot = updater.bot dispatcher = updater.dispatcher diff --git a/bot/helper/telegram_helper/message_utils.py b/bot/helper/telegram_helper/message_utils.py index 9a15264..ad4e6fe 100644 --- a/bot/helper/telegram_helper/message_utils.py +++ b/bot/helper/telegram_helper/message_utils.py @@ -9,31 +9,25 @@ def sendMessage(text: str, bot, update: Update): - try: - return bot.send_message(update.message.chat_id, - reply_to_message_id=update.message.message_id, - text=text, parse_mode='HTMl',timeout=60) - except TimedOut as e: - LOGGER.error(str(e)) - time.sleep(5) #sleep for some time and try request again - return bot.send_message(update.message.chat_id, + return bot.send_message(update.message.chat_id, reply_to_message_id=update.message.message_id, - text=text, parse_mode='HTMl',timeout=60) + text=text, parse_mode='HTMl') def editMessage(text: str, message: Message): try: bot.edit_message_text(text=text, message_id=message.message_id, chat_id=message.chat.id, - parse_mode='HTMl',timeout=60) + parse_mode='HTMl') except TimedOut as e: LOGGER.error(str(e)) + pass def deleteMessage(bot, message: Message): try: bot.delete_message(chat_id=message.chat.id, - message_id=message.message_id,timeout=60) + message_id=message.message_id) except Exception as e: LOGGER.error(str(e)) From 50f04bc40939598996b077c48a87df5eb6298ee3 Mon Sep 17 00:00:00 2001 From: jaskaranSM Date: Mon, 17 Feb 2020 21:09:21 +0530 Subject: [PATCH 22/37] Cleanup drive_list and minor fixes --- .../mirror_utils/upload_utils/gdriveTools.py | 52 ++++++++----------- 1 file changed, 22 insertions(+), 30 deletions(-) diff --git a/bot/helper/mirror_utils/upload_utils/gdriveTools.py b/bot/helper/mirror_utils/upload_utils/gdriveTools.py index b27d940..0286a67 100644 --- a/bot/helper/mirror_utils/upload_utils/gdriveTools.py +++ b/bot/helper/mirror_utils/upload_utils/gdriveTools.py @@ -221,8 +221,8 @@ def authorize(self): # Get credentials credentials = None if not USE_SERVICE_ACCOUNTS: - if os.path.exists(G_DRIVE_TOKEN_FILE): - with open(G_DRIVE_TOKEN_FILE, 'rb') as f: + if os.path.exists(self.__G_DRIVE_TOKEN_FILE): + with open(self.__G_DRIVE_TOKEN_FILE, 'rb') as f: credentials = pickle.load(f) if credentials is None or not credentials.valid: if credentials and credentials.expired and credentials.refresh_token: @@ -234,7 +234,7 @@ def authorize(self): credentials = flow.run_console(port=0) # Save the credentials for the next run - with open(G_DRIVE_TOKEN_FILE, 'wb') as token: + with open(self.__G_DRIVE_TOKEN_FILE, 'wb') as token: pickle.dump(credentials, token) else: LOGGER.info(f"Authorizing with {SERVICE_ACCOUNT_INDEX}.json service account") @@ -246,35 +246,27 @@ def drive_list(self, fileName): msg = "" # Create Search Query for API request. query = f"'{parent_id}' in parents and (name contains '{fileName}')" - page_token = None - count = 0 - while True: - response = self.__service.files().list(supportsTeamDrives=True, + response = self.__service.files().list(supportsTeamDrives=True, includeTeamDriveItems=True, q=query, spaces='drive', - fields='nextPageToken, files(id, name, mimeType, size)', - pageToken=page_token, + pageSize=20, + fields='files(id, name, mimeType, size)', orderBy='modifiedTime desc').execute() - for file in response.get('files', []): - if count >= 20: - break - if file.get( - 'mimeType') == "application/vnd.google-apps.folder": # Detect Whether Current Entity is a Folder or File. - msg += f"⁍ {file.get('name')}" \ - f" (folder)" - if INDEX_URL is not None: - url = f'{INDEX_URL}/{file.get("name")}/' - msg += f' | Index URL' - else: - msg += f"⁍ {file.get('name')} ({get_readable_file_size(int(file.get('size')))})" - if INDEX_URL is not None: - url = f'{INDEX_URL}/{file.get("name")}' - msg += f' | Index URL' - msg += '\n' - count += 1 - page_token = response.get('nextPageToken', None) - if page_token is None or count >= 20: - break + for file in response.get('files', []): + if file.get( + 'mimeType') == "application/vnd.google-apps.folder": # Detect Whether Current Entity is a Folder or File. + msg += f"⁍ {file.get('name')}" \ + f" (folder)" + if INDEX_URL is not None: + url = f'{INDEX_URL}/{file.get("name")}/' + msg += f' | Index URL' + else: + msg += f"⁍ {file.get('name')} ({get_readable_file_size(int(file.get('size')))})" + if INDEX_URL is not None: + url = f'{INDEX_URL}/{file.get("name")}' + msg += f' | Index URL' + msg += '\n' + return msg From 53dc82c33ccfb53afedd34373b65afe6fb7a2c11 Mon Sep 17 00:00:00 2001 From: jaskaranSM Date: Sat, 21 Mar 2020 10:31:24 +0530 Subject: [PATCH 23/37] recursively try again instead of handling switch with backoff - Handling service account switching with exponential backoff implementation limits the number of service accounts to be used. recursively retrying again after switch doesnt effect backoff. --- .../mirror_utils/upload_utils/gdriveTools.py | 22 ++++++++++++++----- 1 file changed, 16 insertions(+), 6 deletions(-) diff --git a/bot/helper/mirror_utils/upload_utils/gdriveTools.py b/bot/helper/mirror_utils/upload_utils/gdriveTools.py index 0286a67..e1df778 100644 --- a/bot/helper/mirror_utils/upload_utils/gdriveTools.py +++ b/bot/helper/mirror_utils/upload_utils/gdriveTools.py @@ -40,6 +40,7 @@ def __init__(self, name=None, listener=None): self.status = None self.updater = None self.name = name + self.service_account_count = len(os.listdir("accounts")) def cancel(self): self.is_cancelled = True @@ -65,6 +66,14 @@ def _on_upload_progress(self): self.uploaded_bytes += chunk_size self.total_time += self.UPDATE_INTERVAL + def switchServiceAccount(self): + global SERVICE_ACCOUNT_INDEX + if SERVICE_ACCOUNT_INDEX == self.service_account_count - 1: + SERVICE_ACCOUNT_INDEX = 0 + SERVICE_ACCOUNT_INDEX += 1 + LOGGER.info(f"Switching to {SERVICE_ACCOUNT_INDEX}.json service account") + self.__service = self.authorize() + @retry(wait=wait_exponential(multiplier=2, min=3, max=6), stop=stop_after_attempt(5), retry=retry_if_exception_type(HttpError), before=before_log(LOGGER, logging.DEBUG)) def __set_permission(self, drive_id): @@ -117,12 +126,13 @@ def upload_file(self, file_path, file_name, mime_type, parent_id): 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': - global SERVICE_ACCOUNT_INDEX - SERVICE_ACCOUNT_INDEX += 1 - LOGGER.info(f"Switching to {SERVICE_ACCOUNT_INDEX}.json service account") - self.__service = self.authorize() - raise err + if reason == 'userRateLimitExceeded' or reason == 'dailyLimitExceeded': + if USE_SERVICE_ACCOUNTS: + self.switchServiceAccount() + LOGGER.info(f"Got: {reason}, Trying Again.") + self.upload_file(file_path, file_name, mime_type, parent_id) + else: + raise err self._file_uploaded_bytes = 0 # Insert new permissions if not IS_TEAM_DRIVE: From 8207520759b9b7e2394821cddac6a1c6e6a6fe9d Mon Sep 17 00:00:00 2001 From: jaskaranSM Date: Sat, 21 Mar 2020 11:42:07 +0530 Subject: [PATCH 24/37] Implement clone feature - clones a public/driveAccountAccessible Drive Link to bot G-Drive --- .../mirror_utils/upload_utils/gdriveTools.py | 80 +++++++++++++++++++ bot/modules/clone.py | 22 +++++ 2 files changed, 102 insertions(+) create mode 100644 bot/modules/clone.py diff --git a/bot/helper/mirror_utils/upload_utils/gdriveTools.py b/bot/helper/mirror_utils/upload_utils/gdriveTools.py index e1df778..f77215d 100644 --- a/bot/helper/mirror_utils/upload_utils/gdriveTools.py +++ b/bot/helper/mirror_utils/upload_utils/gdriveTools.py @@ -26,6 +26,7 @@ def __init__(self, name=None, listener=None): self.__REDIRECT_URI = "urn:ietf:wg:oauth:2.0:oob" self.__G_DRIVE_DIR_MIME_TYPE = "application/vnd.google-apps.folder" self.__G_DRIVE_BASE_DOWNLOAD_URL = "https://drive.google.com/uc?id={}&export=download" + self.__G_DRIVE_DIR_BASE_DOWNLOAD_URL = "https://drive.google.com/drive/folders/{}" self.__listener = listener self.__service = self.authorize() self.__listener = listener @@ -56,6 +57,13 @@ def speed(self): except ZeroDivisionError: return 0 + def parseLink(self,link): + if "folders" in link or "file" in link: + node_id = link.split("?")[0].split("/")[-1].replace("/",'') + else: + node_id = link.split("=")[1].split("&")[0].replace("/",'') + return node_id + @retry(wait=wait_exponential(multiplier=2, min=3, max=6), stop=stop_after_attempt(5), retry=retry_if_exception_type(HttpError), before=before_log(LOGGER, logging.DEBUG)) def _on_upload_progress(self): @@ -191,6 +199,78 @@ def upload(self, file_name: str): LOGGER.info("Deleting downloaded file/folder..") return link + @retry(wait=wait_exponential(multiplier=2, min=3, max=6),stop=stop_after_attempt(5),retry=retry_if_exception_type(HttpError),before=before_log(LOGGER,logging.DEBUG)) + def copyFile(self,file_id,dest_id): + body = { + 'parents': [dest_id] + } + return self.__service.files().copy(supportsAllDrives=True,fileId=file_id,body=body).execute() + + def clone(self,link): + self.transferred_size = 0 + file_id = self.parseLink(link) + msg = "" + LOGGER.info(f"File ID: {file_id}") + try: + meta = self.__service.files().get(fileId=file_id,fields="name,id,mimeType,size").execute() + except Exception as e: + return f"{str(e).replace('>','').replace('<','')}" + if meta.get("mimeType") == self.__G_DRIVE_DIR_MIME_TYPE: + dir_id = self.create_directory(meta.get('name'),parent_id) + try: + result = self.cloneFolder(meta.get('name'),meta.get('name'),meta.get('id'),dir_id) + except Exception as e: + if isinstance(e,RetryError): + LOGGER.info(f"Total Attempts: {e.last_attempt.attempt_number}") + err = e.last_attempt.exception() + else: + err = str(e).replace('>','').replace('<','') + LOGGER.error(err) + return err + msg += f'{meta.get("name")} ({get_readable_file_size(self.transferred_size)})' + else: + file = self.copyFile(meta.get('id'),parent_id) + msg += f'{meta.get("name")} ({get_readable_file_size(int(meta.get("size")))})' + return msg + + + def cloneFolder(self,name,local_path,folder_id,parent_id): + page_token = None + q =f"'{folder_id}' in parents" + files = [] + LOGGER.info(f"Syncing: {local_path}") + new_id = None + while True: + response = self.__service.files().list(q=q, + spaces='drive', + fields='nextPageToken, files(id, name, mimeType,size)', + pageToken=page_token).execute() + for file in response.get('files', []): + files.append(file) + page_token = response.get('nextPageToken', None) + if page_token is None: + break + if len(files) == 0: + return parent_id + for file in files: + if file.get('mimeType') == self.__G_DRIVE_DIR_MIME_TYPE: + file_path = os.path.join(local_path,file.get('name')) + current_dir_id = self.create_directory(file.get('name'),parent_id) + new_id = self.cloneFolder(file.get('name'),file_path,file.get('id'),current_dir_id) + else: + self.transferred_size += int(file.get('size')) + try: + self.copyFile(file.get('id'),parent_id) + new_id = parent_id + except Exception as e: + if isinstance(e,RetryError): + LOGGER.info(f"Total Attempts: {e.last_attempt.attempt_number}") + err = e.last_attempt.exception() + else: + err = e + LOGGER.error(err) + return new_id + @retry(wait=wait_exponential(multiplier=2, min=3, max=6), stop=stop_after_attempt(5), retry=retry_if_exception_type(HttpError), before=before_log(LOGGER, logging.DEBUG)) def create_directory(self, directory_name, parent_id): diff --git a/bot/modules/clone.py b/bot/modules/clone.py new file mode 100644 index 0000000..9107f10 --- /dev/null +++ b/bot/modules/clone.py @@ -0,0 +1,22 @@ +from telegram.ext import CommandHandler, run_async +from bot.helper.mirror_utils.upload_utils.gdriveTools import GoogleDriveHelper +from bot.helper.telegram_helper.message_utils import * +from bot.helper.telegram_helper.filters import CustomFilters +from bot import dispatcher + + +@run_async +def cloneNode(bot,update): + args = update.message.text.split(" ",maxsplit=1) + if len(args) > 1: + link = args[1] + msg = sendMessage(f"Cloning: {link}",bot,update) + gd = GoogleDriveHelper() + result = gd.clone(link) + deleteMessage(bot,msg) + sendMessage(result,bot,update) + else: + sendMessage("Provide G-Drive Shareable Link to Clone.",bot,update) + +clone_handler = CommandHandler('clone',cloneNode,filters=CustomFilters.authorized_chat | CustomFilters.authorized_user) +dispatcher.add_handler(clone_handler) \ No newline at end of file From 58fd47bcc01fed6742df44e1c083309df9666622 Mon Sep 17 00:00:00 2001 From: jaskaranSM Date: Sat, 21 Mar 2020 11:45:21 +0530 Subject: [PATCH 25/37] use BotCommands for clone command trigger --- bot/helper/telegram_helper/bot_commands.py | 1 + bot/modules/clone.py | 3 ++- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/bot/helper/telegram_helper/bot_commands.py b/bot/helper/telegram_helper/bot_commands.py index 918500d..d58c29d 100644 --- a/bot/helper/telegram_helper/bot_commands.py +++ b/bot/helper/telegram_helper/bot_commands.py @@ -13,6 +13,7 @@ def __init__(self): self.StatsCommand = 'stats' self.HelpCommand = 'help' self.LogCommand = 'log' + self.CloneCommand = "clone" BotCommands = _BotCommands() diff --git a/bot/modules/clone.py b/bot/modules/clone.py index 9107f10..fa33233 100644 --- a/bot/modules/clone.py +++ b/bot/modules/clone.py @@ -2,6 +2,7 @@ from bot.helper.mirror_utils.upload_utils.gdriveTools import GoogleDriveHelper from bot.helper.telegram_helper.message_utils import * from bot.helper.telegram_helper.filters import CustomFilters +from bot.helper.telegram_helper.bot_commands import BotCommands from bot import dispatcher @@ -18,5 +19,5 @@ def cloneNode(bot,update): else: sendMessage("Provide G-Drive Shareable Link to Clone.",bot,update) -clone_handler = CommandHandler('clone',cloneNode,filters=CustomFilters.authorized_chat | CustomFilters.authorized_user) +clone_handler = CommandHandler(BotCommands.CloneCommand,cloneNode,filters=CustomFilters.authorized_chat | CustomFilters.authorized_user) dispatcher.add_handler(clone_handler) \ No newline at end of file From 93166cb6b0a4553cf20ba48633f7d7e777b19baf Mon Sep 17 00:00:00 2001 From: jaskaranSM <37726998+jaskaranSM@users.noreply.github.com> Date: Sat, 21 Mar 2020 13:11:15 +0530 Subject: [PATCH 26/37] Import Clone Module --- bot/__main__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bot/__main__.py b/bot/__main__.py index 0c0d2d6..7255746 100644 --- a/bot/__main__.py +++ b/bot/__main__.py @@ -8,7 +8,7 @@ import shutil from .helper.telegram_helper.filters import CustomFilters from bot.helper.telegram_helper.bot_commands import BotCommands -from .modules import authorize, list, cancel_mirror, mirror_status, mirror +from .modules import authorize, list, cancel_mirror, mirror_status, mirror, clone @run_async def stats(bot,update): From e79293ec74624c8a37f4c5922d259eabaaf3f475 Mon Sep 17 00:00:00 2001 From: jaskaranSM Date: Sat, 21 Mar 2020 14:28:38 +0530 Subject: [PATCH 27/37] Implement cancellation of downloads by gid --- bot/helper/ext_utils/bot_utils.py | 9 +++++++ bot/modules/cancel_mirror.py | 44 ++++++++++++++++++++----------- 2 files changed, 37 insertions(+), 16 deletions(-) diff --git a/bot/helper/ext_utils/bot_utils.py b/bot/helper/ext_utils/bot_utils.py index 5f42e6a..2ef75eb 100644 --- a/bot/helper/ext_utils/bot_utils.py +++ b/bot/helper/ext_utils/bot_utils.py @@ -54,6 +54,14 @@ def get_readable_file_size(size_in_bytes) -> str: except IndexError: return 'File too large' +def getDownloadByGid(gid): + with download_dict_lock: + for dl in download_dict.values(): + if dl.status() == MirrorStatus.STATUS_DOWNLOADING: + if dl.download().gid == gid: + return dl + return None + def get_progress_bar_string(status): completed = status.processed_bytes() / 8 @@ -95,6 +103,7 @@ def get_readable_message(): if hasattr(download,'is_torrent'): msg += f"| P: {download.download().connections} " \ f"| S: {download.download().num_seeders}" + msg += f"\nGID: {download.download().gid}" msg += "\n\n" return msg diff --git a/bot/modules/cancel_mirror.py b/bot/modules/cancel_mirror.py index b69079d..fd8eae4 100644 --- a/bot/modules/cancel_mirror.py +++ b/bot/modules/cancel_mirror.py @@ -5,20 +5,36 @@ from bot.helper.ext_utils.fs_utils import clean_download from bot.helper.telegram_helper.bot_commands import BotCommands from time import sleep +from bot.helper.ext_utils.bot_utils import getDownloadByGid @run_async def cancel_mirror(bot,update): - mirror_message = update.message.reply_to_message - with download_dict_lock: - keys = download_dict.keys() - dl = download_dict[mirror_message.message_id] - if mirror_message is None or mirror_message.message_id not in keys: - if '/mirror' in mirror_message.text or '/tarmirror' in mirror_message.text: - msg = 'Message has already been cancelled' - else: - msg = 'Please reply to the /mirror message which was used to start the download to cancel it!' - sendMessage(msg, bot, update) - return + args = update.message.text.split(" ",maxsplit=1) + mirror_message = None + if len(args) > 1: + gid = args[1] + dl = getDownloadByGid(gid) + if not dl: + sendMessage(f"GID: {gid} not found.",bot,update) + return + with download_dict_lock: + keys = list(download_dict.keys()) + mirror_message = dl._listener.message + elif update.message.reply_to_message: + mirror_message = update.message.reply_to_message + with download_dict_lock: + keys = list(download_dict.keys()) + dl = download_dict[mirror_message.message_id] + if len(args) == 1: + if mirror_message is None or mirror_message.message_id not in keys: + if BotCommands.MirrorCommand in mirror_message.text or BotCommands.TarMirrorCommand in mirror_message.text: + msg = "Mirror already have been cancelled" + sendMessage(msg,bot,update) + return + else: + msg = "Please reply to the /mirror message which was used to start the download or /cancel gid to cancel it!" + sendMessage(msg,bot,update) + return if dl.status() == "Uploading": sendMessage("Upload in Progress, Don't Cancel it.", bot, update) return @@ -31,13 +47,9 @@ def cancel_mirror(bot,update): downloads = aria2.get_downloads(download.followed_by_ids) aria2.pause(downloads) aria2.pause([download]) - - elif dl.status() == "Uploading": - sendMessage("Upload in Progress, Dont Cancel it.",bot,update) - return else: dl._listener.onDownloadError("Download stopped by user!") - sleep(1) #Wait a Second For Aria2 To free Resources. + sleep(1) # Wait a Second For Aria2 To free Resources. clean_download(f'{DOWNLOAD_DIR}{mirror_message.message_id}/') From 1e135a1bbce29d549cc078e750099c71af5b5c8e Mon Sep 17 00:00:00 2001 From: jaskaranSM <37726998+jaskaranSM@users.noreply.github.com> Date: Mon, 23 Mar 2020 16:53:13 +0530 Subject: [PATCH 28/37] Enable TeamDrive support for clone --- bot/helper/mirror_utils/upload_utils/gdriveTools.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bot/helper/mirror_utils/upload_utils/gdriveTools.py b/bot/helper/mirror_utils/upload_utils/gdriveTools.py index f77215d..27bbcd3 100644 --- a/bot/helper/mirror_utils/upload_utils/gdriveTools.py +++ b/bot/helper/mirror_utils/upload_utils/gdriveTools.py @@ -212,7 +212,7 @@ def clone(self,link): msg = "" LOGGER.info(f"File ID: {file_id}") try: - meta = self.__service.files().get(fileId=file_id,fields="name,id,mimeType,size").execute() + meta = self.__service.files().get(supportsAllDrives=True,fileId=file_id,fields="name,id,mimeType,size").execute() except Exception as e: return f"{str(e).replace('>','').replace('<','')}" if meta.get("mimeType") == self.__G_DRIVE_DIR_MIME_TYPE: From 9de2a3f8abd5818bf5e94d71267c55c0d6d478e5 Mon Sep 17 00:00:00 2001 From: jaskaranSM <37726998+jaskaranSM@users.noreply.github.com> Date: Mon, 23 Mar 2020 16:59:09 +0530 Subject: [PATCH 29/37] TeamDrive List fix --- bot/helper/mirror_utils/upload_utils/gdriveTools.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/bot/helper/mirror_utils/upload_utils/gdriveTools.py b/bot/helper/mirror_utils/upload_utils/gdriveTools.py index 27bbcd3..39c9786 100644 --- a/bot/helper/mirror_utils/upload_utils/gdriveTools.py +++ b/bot/helper/mirror_utils/upload_utils/gdriveTools.py @@ -241,7 +241,9 @@ def cloneFolder(self,name,local_path,folder_id,parent_id): LOGGER.info(f"Syncing: {local_path}") new_id = None while True: - response = self.__service.files().list(q=q, + response = self.__service.files().list(supportsTeamDrives=True, + includeTeamDriveItems=True, + q=q, spaces='drive', fields='nextPageToken, files(id, name, mimeType,size)', pageToken=page_token).execute() From 47cfaf7d321431825a8015fe9daf10e7e6218527 Mon Sep 17 00:00:00 2001 From: jaskaranSM Date: Tue, 24 Mar 2020 09:59:26 +0530 Subject: [PATCH 30/37] Use Service Accounts for Cloning if availble. - increase default timeout for network socket interface --- bot/__init__.py | 3 +++ .../mirror_utils/upload_utils/gdriveTools.py | 15 ++++++++++++++- 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/bot/__init__.py b/bot/__init__.py index 3bc56e3..900afb4 100644 --- a/bot/__init__.py +++ b/bot/__init__.py @@ -6,6 +6,9 @@ import aria2p import telegram.ext as tg from dotenv import load_dotenv +import socket + +socket.setdefaulttimeout(600) botStartTime = time.time() if os.path.exists('log.txt'): diff --git a/bot/helper/mirror_utils/upload_utils/gdriveTools.py b/bot/helper/mirror_utils/upload_utils/gdriveTools.py index 39c9786..57f4c82 100644 --- a/bot/helper/mirror_utils/upload_utils/gdriveTools.py +++ b/bot/helper/mirror_utils/upload_utils/gdriveTools.py @@ -204,7 +204,20 @@ def copyFile(self,file_id,dest_id): body = { 'parents': [dest_id] } - return self.__service.files().copy(supportsAllDrives=True,fileId=file_id,body=body).execute() + try: + res = self.__service.files().copy(supportsAllDrives=True,fileId=file_id,body=body).execute() + return res + 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: + self.switchServiceAccount() + LOGGER.info(f"Got: {reason}, Trying Again.") + self.copyFile(file_id,dest_id) + else: + raise err + def clone(self,link): self.transferred_size = 0 From 8c3284d04b7b5f811164b573d7d2beff9d75d30a Mon Sep 17 00:00:00 2001 From: jaskaranSM Date: Fri, 27 Mar 2020 19:57:53 +0530 Subject: [PATCH 31/37] Added Youtube-dl support --- .../youtube_dl_download_helper.py | 104 ++++++++++++++++++ .../youtube_dl_download_status.py | 54 +++++++++ bot/modules/cancel_mirror.py | 12 +- bot/modules/mirror.py | 13 ++- requirements.txt | 3 +- 5 files changed, 179 insertions(+), 7 deletions(-) create mode 100644 bot/helper/mirror_utils/download_utils/youtube_dl_download_helper.py create mode 100644 bot/helper/mirror_utils/status_utils/youtube_dl_download_status.py diff --git a/bot/helper/mirror_utils/download_utils/youtube_dl_download_helper.py b/bot/helper/mirror_utils/download_utils/youtube_dl_download_helper.py new file mode 100644 index 0000000..96d3fe6 --- /dev/null +++ b/bot/helper/mirror_utils/download_utils/youtube_dl_download_helper.py @@ -0,0 +1,104 @@ +from .download_helper import DownloadHelper +import time +from youtube_dl import YoutubeDL +import threading +from bot import LOGGER,download_dict_lock,download_dict,DOWNLOAD_DIR +from ..status_utils.youtube_dl_download_status import YoutubeDLDownloadStatus + +class YoutubeDLHelper(DownloadHelper): + def __init__(self, listener): + super().__init__() + self.__name = "" + self.__start_time = time.time() + self.__listener = listener + self.__gid = "" + self.opts = { + 'format': 'bestaudio/best', + 'progress_hooks':[self.__onDownloadProgress], + 'outtmpl': f"{DOWNLOAD_DIR}{self.__listener.uid}/%(title)s.%(ext)s" + } + self.ydl = YoutubeDL(self.opts) + self.__download_speed = 0 + self.download_speed_readable = "" + self.downloaded_bytes = 0 + self.size = 0 + self.is_playlist = False + self.last_downloaded = 0 + self.is_cancelled = False + self.__resource_lock = threading.RLock() + + @property + def download_speed(self): + with self.__resource_lock: + return self.__download_speed + + @property + def gid(self): + with self.__resource_lock: + return self.__gid + + def __onDownloadProgress(self,d): + if self.is_cancelled: + raise ValueError("Cancelling Download..") + if d['status'] == "finished": + if self.is_playlist: + self.last_downloaded = 0 + if self.downloaded_bytes == self.size: + self.__onDownloadComplete() + else: + self.__onDownloadComplete() + elif d['status'] == "downloading": + with self.__resource_lock: + self.progress = self.downloaded_bytes / self.size * 100 + self.__download_speed = d['speed'] + if self.is_playlist: + chunk_size = self.size * self.progress - self.last_downloaded + self.last_downloaded = self.size * self.progress + self.downloaded_bytes += chunk_size + else: + self.download_speed_readable = d['_speed_str'] + self.downloaded_bytes = d['downloaded_bytes'] + + def __onDownloadStart(self): + with download_dict_lock: + download_dict[self.__listener.uid] = YoutubeDLDownloadStatus(self,self.__listener.uid) + + def __onDownloadComplete(self): + self.__listener.onDownloadComplete() + + def __onDownloadError(self,error): + self.__listener.onDownloadError(error) + + def extractMetaData(self,link): + result = self.ydl.extract_info(link,download=False) + if 'entries' in result: + video = result['entries'][0] + for v in result['entries']: + self.size += int(v['filesize']) + self.name = result.get('title') + self.vid_id = video.get('id') + self.is_playlist = True + self.opts['outtmpl'] = f"{DOWNLOAD_DIR}{self.__listener.uid}/%(playlist)s/%(title)s.%(ext)s" + self.ydl = YoutubeDL(self.opts) + else: + video = result + self.size = int(video.get('filesize')) + self.name = video.get('title') + self.vid_id = video.get('id') + return video + + def __download(self,link): + try: + self.ydl.download([link],) + except ValueError: + LOGGER.info("Download Cancelled by User!") + self.__onDownloadError("Download Cancelled by User!") + + def add_download(self,link): + LOGGER.info(f"Downloading with YT-DL: {link}") + self.__gid = f"{self.vid_id}{self.__listener.uid}" + threading.Thread(target=self.__download,args=(link,)).start() + self.__onDownloadStart() + + def cancel_download(self): + self.is_cancelled = True \ No newline at end of file diff --git a/bot/helper/mirror_utils/status_utils/youtube_dl_download_status.py b/bot/helper/mirror_utils/status_utils/youtube_dl_download_status.py new file mode 100644 index 0000000..16b3bff --- /dev/null +++ b/bot/helper/mirror_utils/status_utils/youtube_dl_download_status.py @@ -0,0 +1,54 @@ +from bot import DOWNLOAD_DIR +from bot.helper.ext_utils.bot_utils import MirrorStatus, get_readable_file_size, get_readable_time +from .status import Status + +class YoutubeDLDownloadStatus(Status): + def __init__(self, obj, uid): + self.obj = obj + self.uid = uid + + def gid(self): + return self.obj.gid + + def path(self): + return f"{DOWNLOAD_DIR}{self.uid}" + + def processed_bytes(self): + return self.obj.downloaded_bytes + + def size_raw(self): + return self.obj.size + + def size(self): + return get_readable_file_size(self.size_raw()) + + def status(self): + return MirrorStatus.STATUS_DOWNLOADING + + def name(self): + return self.obj.name + + def progress_raw(self): + return self.obj.progress + + def progress(self): + return f'{round(self.progress_raw(), 2)}%' + + def speed_raw(self): + """ + :return: Download speed in Bytes/Seconds + """ + return self.obj.download_speed + + def speed(self): + return f'{get_readable_file_size(self.speed_raw())}/s' + + def eta(self): + try: + seconds = (self.size_raw() - self.processed_bytes()) / self.speed_raw() + return f'{get_readable_time(seconds)}' + except ZeroDivisionError: + return '-' + + def download(self): + return self.obj diff --git a/bot/modules/cancel_mirror.py b/bot/modules/cancel_mirror.py index fd8eae4..2505fd4 100644 --- a/bot/modules/cancel_mirror.py +++ b/bot/modules/cancel_mirror.py @@ -6,6 +6,7 @@ from bot.helper.telegram_helper.bot_commands import BotCommands from time import sleep from bot.helper.ext_utils.bot_utils import getDownloadByGid +from ..helper.mirror_utils.download_utils.youtube_dl_download_helper import YoutubeDLHelper @run_async def cancel_mirror(bot,update): @@ -43,10 +44,13 @@ def cancel_mirror(bot,update): return elif dl.status() != "Queued": download = dl.download() - if len(download.followed_by_ids) != 0: - downloads = aria2.get_downloads(download.followed_by_ids) - aria2.pause(downloads) - aria2.pause([download]) + if isinstance(download,YoutubeDLHelper): + download.cancel_download() + else: + if len(download.followed_by_ids) != 0: + downloads = aria2.get_downloads(download.followed_by_ids) + aria2.pause(downloads) + aria2.pause([download]) else: dl._listener.onDownloadError("Download stopped by user!") sleep(1) # Wait a Second For Aria2 To free Resources. diff --git a/bot/modules/mirror.py b/bot/modules/mirror.py index 4243118..f511963 100644 --- a/bot/modules/mirror.py +++ b/bot/modules/mirror.py @@ -11,6 +11,7 @@ from bot.helper.ext_utils.bot_utils import setInterval from bot.helper.telegram_helper.filters import CustomFilters from bot.helper.telegram_helper.bot_commands import BotCommands +from bot.helper.mirror_utils.download_utils.youtube_dl_download_helper import YoutubeDLHelper import pathlib import os @@ -148,8 +149,16 @@ def _mirror(bot, update, isTar=False): sendMessage('No download source provided', bot, update) return listener = MirrorListener(bot, update, isTar, tag) - aria = aria2_download.AriaDownloadHelper(listener) - aria.add_download(link, f'{DOWNLOAD_DIR}/{listener.uid}/') + try: + ydl = YoutubeDLHelper(listener) + sup_link = ydl.extractMetaData(link) + except Exception as e: + sup_link = None + if sup_link: + ydl.add_download(link) + else: + aria = aria2_download.AriaDownloadHelper(listener) + aria.add_download(link, f'{DOWNLOAD_DIR}/{listener.uid}/') sendStatusMessage(update, bot) if len(Interval) == 0: Interval.append(setInterval(DOWNLOAD_STATUS_UPDATE_INTERVAL, update_all_messages)) diff --git a/requirements.txt b/requirements.txt index a917da1..151622b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,4 +4,5 @@ google-auth-httplib2>=0.0.3,<0.1.0 google-auth-oauthlib>=0.4.1,<0.10.0 aria2p>=0.3.0,<0.10.0 python-dotenv>=0.10 -tenacity>=6.0.0 +tenacity>=6.0.0 +youtube-dl From d5408c488673a39ff4f9b0812c14dd579deb6016 Mon Sep 17 00:00:00 2001 From: lzzy12 Date: Sat, 28 Mar 2020 13:47:19 +0530 Subject: [PATCH 32/37] Fixed yt-dl support Signed-off-by: lzzy12 --- .../youtube_dl_download_helper.py | 48 ++++++++++--------- .../youtube_dl_download_status.py | 1 + bot/modules/mirror.py | 4 +- 3 files changed, 29 insertions(+), 24 deletions(-) diff --git a/bot/helper/mirror_utils/download_utils/youtube_dl_download_helper.py b/bot/helper/mirror_utils/download_utils/youtube_dl_download_helper.py index 96d3fe6..dfc12cb 100644 --- a/bot/helper/mirror_utils/download_utils/youtube_dl_download_helper.py +++ b/bot/helper/mirror_utils/download_utils/youtube_dl_download_helper.py @@ -2,9 +2,10 @@ import time from youtube_dl import YoutubeDL import threading -from bot import LOGGER,download_dict_lock,download_dict,DOWNLOAD_DIR +from bot import LOGGER, download_dict_lock, download_dict, DOWNLOAD_DIR from ..status_utils.youtube_dl_download_status import YoutubeDLDownloadStatus + class YoutubeDLHelper(DownloadHelper): def __init__(self, listener): super().__init__() @@ -13,18 +14,18 @@ def __init__(self, listener): self.__listener = listener self.__gid = "" self.opts = { - 'format': 'bestaudio/best', - 'progress_hooks':[self.__onDownloadProgress], - 'outtmpl': f"{DOWNLOAD_DIR}{self.__listener.uid}/%(title)s.%(ext)s" + 'format': 'bestaudio/best', + 'progress_hooks': [self.__onDownloadProgress], } self.ydl = YoutubeDL(self.opts) self.__download_speed = 0 - self.download_speed_readable = "" + self.download_speed_readable = '' self.downloaded_bytes = 0 self.size = 0 self.is_playlist = False self.last_downloaded = 0 self.is_cancelled = False + self.vid_id = '' self.__resource_lock = threading.RLock() @property @@ -37,7 +38,7 @@ def gid(self): with self.__resource_lock: return self.__gid - def __onDownloadProgress(self,d): + def __onDownloadProgress(self, d): if self.is_cancelled: raise ValueError("Cancelling Download..") if d['status'] == "finished": @@ -45,40 +46,41 @@ def __onDownloadProgress(self,d): self.last_downloaded = 0 if self.downloaded_bytes == self.size: self.__onDownloadComplete() - else: + else: self.__onDownloadComplete() elif d['status'] == "downloading": with self.__resource_lock: - self.progress = self.downloaded_bytes / self.size * 100 self.__download_speed = d['speed'] if self.is_playlist: - chunk_size = self.size * self.progress - self.last_downloaded - self.last_downloaded = self.size * self.progress + progress = d['downloaded_bytes'] / d['total_bytes'] + chunk_size = d['downloaded_bytes'] - self.last_downloaded + self.last_downloaded = d['total_bytes'] * progress self.downloaded_bytes += chunk_size + self.progress = (self.downloaded_bytes / self.size) * 100 else: self.download_speed_readable = d['_speed_str'] self.downloaded_bytes = d['downloaded_bytes'] - + def __onDownloadStart(self): with download_dict_lock: - download_dict[self.__listener.uid] = YoutubeDLDownloadStatus(self,self.__listener.uid) + download_dict[self.__listener.uid] = YoutubeDLDownloadStatus(self, self.__listener.uid) def __onDownloadComplete(self): self.__listener.onDownloadComplete() - def __onDownloadError(self,error): + def __onDownloadError(self, error): self.__listener.onDownloadError(error) - def extractMetaData(self,link): - result = self.ydl.extract_info(link,download=False) + def extractMetaData(self, link): + result = self.ydl.extract_info(link, download=False) if 'entries' in result: video = result['entries'][0] for v in result['entries']: - self.size += int(v['filesize']) + self.size += float(v['filesize']) self.name = result.get('title') self.vid_id = video.get('id') self.is_playlist = True - self.opts['outtmpl'] = f"{DOWNLOAD_DIR}{self.__listener.uid}/%(playlist)s/%(title)s.%(ext)s" + self.opts['o'] = f"{DOWNLOAD_DIR}{self.__listener.uid}/%(playlist)s/%(title)s.%(ext)s" self.ydl = YoutubeDL(self.opts) else: video = result @@ -87,18 +89,20 @@ def extractMetaData(self,link): self.vid_id = video.get('id') return video - def __download(self,link): + def __download(self, link): try: - self.ydl.download([link],) + self.ydl.download([link], ) + self.__onDownloadComplete() except ValueError: LOGGER.info("Download Cancelled by User!") self.__onDownloadError("Download Cancelled by User!") - def add_download(self,link): + def add_download(self, link, path): LOGGER.info(f"Downloading with YT-DL: {link}") self.__gid = f"{self.vid_id}{self.__listener.uid}" - threading.Thread(target=self.__download,args=(link,)).start() + self.opts['o'] = f"{path}/%(title)s/%(title)s.%(ext)s" self.__onDownloadStart() + threading.Thread(target=self.__download, args=(link,)).start() def cancel_download(self): - self.is_cancelled = True \ No newline at end of file + self.is_cancelled = True diff --git a/bot/helper/mirror_utils/status_utils/youtube_dl_download_status.py b/bot/helper/mirror_utils/status_utils/youtube_dl_download_status.py index 16b3bff..51f0582 100644 --- a/bot/helper/mirror_utils/status_utils/youtube_dl_download_status.py +++ b/bot/helper/mirror_utils/status_utils/youtube_dl_download_status.py @@ -2,6 +2,7 @@ from bot.helper.ext_utils.bot_utils import MirrorStatus, get_readable_file_size, get_readable_time from .status import Status + class YoutubeDLDownloadStatus(Status): def __init__(self, obj, uid): self.obj = obj diff --git a/bot/modules/mirror.py b/bot/modules/mirror.py index f511963..2e5d122 100644 --- a/bot/modules/mirror.py +++ b/bot/modules/mirror.py @@ -149,13 +149,13 @@ def _mirror(bot, update, isTar=False): sendMessage('No download source provided', bot, update) return listener = MirrorListener(bot, update, isTar, tag) + ydl = YoutubeDLHelper(listener) try: - ydl = YoutubeDLHelper(listener) sup_link = ydl.extractMetaData(link) except Exception as e: sup_link = None if sup_link: - ydl.add_download(link) + ydl.add_download(link, f'{DOWNLOAD_DIR}/{listener.uid}') else: aria = aria2_download.AriaDownloadHelper(listener) aria.add_download(link, f'{DOWNLOAD_DIR}/{listener.uid}/') From 728a5236d6b49197ef64310c5b8b3197371f5c47 Mon Sep 17 00:00:00 2001 From: lzzy12 Date: Sat, 28 Mar 2020 14:12:32 +0530 Subject: [PATCH 33/37] Fix uploading for YT-DL Signed-off-by: lzzy12 --- .../download_utils/youtube_dl_download_helper.py | 16 +++++++--------- bot/modules/mirror.py | 2 +- 2 files changed, 8 insertions(+), 10 deletions(-) diff --git a/bot/helper/mirror_utils/download_utils/youtube_dl_download_helper.py b/bot/helper/mirror_utils/download_utils/youtube_dl_download_helper.py index dfc12cb..2a590e5 100644 --- a/bot/helper/mirror_utils/download_utils/youtube_dl_download_helper.py +++ b/bot/helper/mirror_utils/download_utils/youtube_dl_download_helper.py @@ -44,10 +44,6 @@ def __onDownloadProgress(self, d): if d['status'] == "finished": if self.is_playlist: self.last_downloaded = 0 - if self.downloaded_bytes == self.size: - self.__onDownloadComplete() - else: - self.__onDownloadComplete() elif d['status'] == "downloading": with self.__resource_lock: self.__download_speed = d['speed'] @@ -77,15 +73,13 @@ def extractMetaData(self, link): video = result['entries'][0] for v in result['entries']: self.size += float(v['filesize']) - self.name = result.get('title') + self.name = video.get('playlist_title') self.vid_id = video.get('id') self.is_playlist = True - self.opts['o'] = f"{DOWNLOAD_DIR}{self.__listener.uid}/%(playlist)s/%(title)s.%(ext)s" - self.ydl = YoutubeDL(self.opts) else: video = result self.size = int(video.get('filesize')) - self.name = video.get('title') + self.name = f"{video.get('title')}.{video.get('ext')}" self.vid_id = video.get('id') return video @@ -100,7 +94,11 @@ def __download(self, link): def add_download(self, link, path): LOGGER.info(f"Downloading with YT-DL: {link}") self.__gid = f"{self.vid_id}{self.__listener.uid}" - self.opts['o'] = f"{path}/%(title)s/%(title)s.%(ext)s" + if not self.is_playlist: + self.opts['outtmpl'] = f"{path}/%(title)s.%(ext)s" + else: + self.opts['outtmpl'] = f"{path}/%(playlist_title)s/%(title)s.%(ext)s" + self.ydl = YoutubeDL(self.opts) self.__onDownloadStart() threading.Thread(target=self.__download, args=(link,)).start() diff --git a/bot/modules/mirror.py b/bot/modules/mirror.py index 2e5d122..2af83e1 100644 --- a/bot/modules/mirror.py +++ b/bot/modules/mirror.py @@ -155,7 +155,7 @@ def _mirror(bot, update, isTar=False): except Exception as e: sup_link = None if sup_link: - ydl.add_download(link, f'{DOWNLOAD_DIR}/{listener.uid}') + ydl.add_download(link, f'{DOWNLOAD_DIR}{listener.uid}') else: aria = aria2_download.AriaDownloadHelper(listener) aria.add_download(link, f'{DOWNLOAD_DIR}/{listener.uid}/') From cb376a059166280424d2a97e4c00573647004226 Mon Sep 17 00:00:00 2001 From: jaskaranSM Date: Sat, 28 Mar 2020 17:37:17 +0530 Subject: [PATCH 34/37] youtube-dl: fixed downloading from sites that provide encrypted content (m3u8) --- bot/helper/ext_utils/fs_utils.py | 10 ++++++++++ .../download_utils/youtube_dl_download_helper.py | 6 ++++-- bot/modules/mirror.py | 8 +++++--- 3 files changed, 19 insertions(+), 5 deletions(-) diff --git a/bot/helper/ext_utils/fs_utils.py b/bot/helper/ext_utils/fs_utils.py index 36955f1..dd534b7 100644 --- a/bot/helper/ext_utils/fs_utils.py +++ b/bot/helper/ext_utils/fs_utils.py @@ -30,6 +30,16 @@ def exit_clean_up(signal, frame): LOGGER.warning("Force Exiting before the cleanup finishes!") sys.exit(1) +def get_path_size(path): + if os.path.isfile(path): + return os.path.getsize(path) + total_size = 0 + for root, dirs, files in os.walk(path): + for f in files: + abs_path = os.path.join(root, f) + total_size += os.path.getsize(abs_path) + return total_size + def tar(org_path): tar_path = org_path + ".tar" diff --git a/bot/helper/mirror_utils/download_utils/youtube_dl_download_helper.py b/bot/helper/mirror_utils/download_utils/youtube_dl_download_helper.py index 2a590e5..334858d 100644 --- a/bot/helper/mirror_utils/download_utils/youtube_dl_download_helper.py +++ b/bot/helper/mirror_utils/download_utils/youtube_dl_download_helper.py @@ -72,13 +72,15 @@ def extractMetaData(self, link): if 'entries' in result: video = result['entries'][0] for v in result['entries']: - self.size += float(v['filesize']) + if v.get('filesize'): + self.size += float(v['filesize']) self.name = video.get('playlist_title') self.vid_id = video.get('id') self.is_playlist = True else: video = result - self.size = int(video.get('filesize')) + if video.get('filesize'): + self.size = int(video.get('filesize')) self.name = f"{video.get('title')}.{video.get('ext')}" self.vid_id = video.get('id') return video diff --git a/bot/modules/mirror.py b/bot/modules/mirror.py index 2af83e1..94ea9ba 100644 --- a/bot/modules/mirror.py +++ b/bot/modules/mirror.py @@ -54,10 +54,12 @@ def onDownloadComplete(self): else: path = f'{DOWNLOAD_DIR}{self.uid}/{download_dict[self.uid].name()}' up_name = pathlib.PurePath(path).name + LOGGER.info(f"Upload Name : {up_name}") + drive = gdriveTools.GoogleDriveHelper(up_name, self) + if size == 0: + size = fs_utils.get_path_size(m_path) + upload_status = UploadStatus(drive, size, self.uid) with download_dict_lock: - LOGGER.info(f"Upload Name : {up_name}") - drive = gdriveTools.GoogleDriveHelper(up_name, self) - upload_status = UploadStatus(drive, size, self.uid) download_dict[self.uid] = upload_status update_all_messages() drive.upload(up_name) From 16e7dc795c3f7aeaa59ea09848fc60032c3ef99d Mon Sep 17 00:00:00 2001 From: jaskaranSM Date: Sun, 29 Mar 2020 13:52:41 +0530 Subject: [PATCH 35/37] fixed direct link mirroring --- .../mirror_utils/download_utils/youtube_dl_download_helper.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/bot/helper/mirror_utils/download_utils/youtube_dl_download_helper.py b/bot/helper/mirror_utils/download_utils/youtube_dl_download_helper.py index 334858d..e7dfcc9 100644 --- a/bot/helper/mirror_utils/download_utils/youtube_dl_download_helper.py +++ b/bot/helper/mirror_utils/download_utils/youtube_dl_download_helper.py @@ -69,6 +69,8 @@ def __onDownloadError(self, error): def extractMetaData(self, link): result = self.ydl.extract_info(link, download=False) + if result.get('direct'): + return None if 'entries' in result: video = result['entries'][0] for v in result['entries']: From 4ac74d3e0665bdab6ff816fb7d7ca61fac7624da Mon Sep 17 00:00:00 2001 From: jaskaranSM Date: Thu, 16 Apr 2020 12:13:08 +0530 Subject: [PATCH 36/37] switch to context based callbacks --- bot/__init__.py | 2 +- bot/__main__.py | 18 +++++++++--------- .../mirror_utils/upload_utils/gdriveTools.py | 3 ++- bot/modules/authorize.py | 8 ++++---- bot/modules/cancel_mirror.py | 16 ++++++++-------- bot/modules/clone.py | 8 ++++---- bot/modules/list.py | 8 ++++---- bot/modules/mirror.py | 8 ++++---- bot/modules/mirror_status.py | 8 ++++---- requirements.txt | 2 +- 10 files changed, 41 insertions(+), 40 deletions(-) diff --git a/bot/__init__.py b/bot/__init__.py index 900afb4..e2ca270 100644 --- a/bot/__init__.py +++ b/bot/__init__.py @@ -100,6 +100,6 @@ def getConfig(name: str): except KeyError: USE_SERVICE_ACCOUNTS = False -updater = tg.Updater(token=BOT_TOKEN) +updater = tg.Updater(token=BOT_TOKEN,use_context=True) bot = updater.bot dispatcher = updater.dispatcher diff --git a/bot/__main__.py b/bot/__main__.py index 7255746..0c5c2a4 100644 --- a/bot/__main__.py +++ b/bot/__main__.py @@ -11,7 +11,7 @@ from .modules import authorize, list, cancel_mirror, mirror_status, mirror, clone @run_async -def stats(bot,update): +def stats(update,context): currentTime = get_readable_time((time.time() - botStartTime)) total, used, free = shutil.disk_usage('.') total = get_readable_file_size(total) @@ -21,27 +21,27 @@ def stats(bot,update): f'Total disk space: {total}\n' \ f'Used: {used}\n' \ f'Free: {free}' - sendMessage(stats, bot, update) + sendMessage(stats, context.bot, update) @run_async -def start(bot,update): +def start(update,context): sendMessage("This is a bot which can mirror all your links to Google drive!\n" - "Type /help to get a list of available commands", bot, update) + "Type /help to get a list of available commands", context.bot, update) @run_async -def ping(bot,update): +def ping(update,context): start_time = int(round(time.time() * 1000)) - reply = sendMessage("Starting Ping", bot, update) + reply = sendMessage("Starting Ping", context.bot, update) end_time = int(round(time.time()*1000)) editMessage(f'{end_time - start_time} ms',reply) @run_async -def log(bot,update): - sendLogFile(bot, update) +def log(update,context): + sendLogFile(context.bot, update) @run_async @@ -66,7 +66,7 @@ def bot_help(bot,update): /{BotCommands.LogCommand}: Get a log file of the bot. Handy for getting crash reports ''' - sendMessage(help_string, bot, update) + sendMessage(help_string, context.bot, update) def main(): diff --git a/bot/helper/mirror_utils/upload_utils/gdriveTools.py b/bot/helper/mirror_utils/upload_utils/gdriveTools.py index 57f4c82..bea9d4f 100644 --- a/bot/helper/mirror_utils/upload_utils/gdriveTools.py +++ b/bot/helper/mirror_utils/upload_utils/gdriveTools.py @@ -41,7 +41,6 @@ def __init__(self, name=None, listener=None): self.status = None self.updater = None self.name = name - self.service_account_count = len(os.listdir("accounts")) def cancel(self): self.is_cancelled = True @@ -151,6 +150,8 @@ def upload_file(self, file_path, file_name, mime_type, parent_id): return download_url def upload(self, file_name: str): + if USE_SERVICE_ACCOUNTS: + self.service_account_count = len(os.listdir("accounts")) self.__listener.onUploadStarted() file_dir = f"{DOWNLOAD_DIR}{self.__listener.message.message_id}" file_path = f"{file_dir}/{file_name}" diff --git a/bot/modules/authorize.py b/bot/modules/authorize.py index de254f0..fedb4d1 100644 --- a/bot/modules/authorize.py +++ b/bot/modules/authorize.py @@ -9,7 +9,7 @@ @run_async -def authorize(bot, update): +def authorize(update,context): reply_message = update.message.reply_to_message msg = '' with open('authorized_chats.txt', 'a') as file: @@ -31,11 +31,11 @@ def authorize(bot, update): msg = 'Person Authorized to use the bot!' else: msg = 'Person already authorized' - sendMessage(msg, bot, update) + sendMessage(msg, context.bot, update) @run_async -def unauthorize(bot,update): +def unauthorize(update,context): reply_message = update.message.reply_to_message if reply_message is None: # Trying to unauthorize a chat @@ -57,7 +57,7 @@ def unauthorize(bot,update): file.truncate(0) for i in AUTHORIZED_CHATS: file.write(f'{i}\n') - sendMessage(msg, bot, update) + sendMessage(msg, context.bot, update) authorize_handler = CommandHandler(command=BotCommands.AuthorizeCommand, callback=authorize, diff --git a/bot/modules/cancel_mirror.py b/bot/modules/cancel_mirror.py index 2505fd4..6d7a968 100644 --- a/bot/modules/cancel_mirror.py +++ b/bot/modules/cancel_mirror.py @@ -9,14 +9,14 @@ from ..helper.mirror_utils.download_utils.youtube_dl_download_helper import YoutubeDLHelper @run_async -def cancel_mirror(bot,update): +def cancel_mirror(update,context): args = update.message.text.split(" ",maxsplit=1) mirror_message = None if len(args) > 1: gid = args[1] dl = getDownloadByGid(gid) if not dl: - sendMessage(f"GID: {gid} not found.",bot,update) + sendMessage(f"GID: {gid} not found.",context.bot,update) return with download_dict_lock: keys = list(download_dict.keys()) @@ -30,17 +30,17 @@ def cancel_mirror(bot,update): if mirror_message is None or mirror_message.message_id not in keys: if BotCommands.MirrorCommand in mirror_message.text or BotCommands.TarMirrorCommand in mirror_message.text: msg = "Mirror already have been cancelled" - sendMessage(msg,bot,update) + sendMessage(msg,context.bot,update) return else: msg = "Please reply to the /mirror message which was used to start the download or /cancel gid to cancel it!" - sendMessage(msg,bot,update) + sendMessage(msg,context.bot,update) return if dl.status() == "Uploading": - sendMessage("Upload in Progress, Don't Cancel it.", bot, update) + sendMessage("Upload in Progress, Don't Cancel it.", context.bot, update) return elif dl.status() == "Archiving": - sendMessage("Archival in Progress, Don't Cancel it.", bot, update) + sendMessage("Archival in Progress, Don't Cancel it.", context.bot, update) return elif dl.status() != "Queued": download = dl.download() @@ -58,7 +58,7 @@ def cancel_mirror(bot,update): @run_async -def cancel_all(update, bot): +def cancel_all(update, context): with download_dict_lock: count = 0 for dlDetails in list(download_dict.values()): @@ -70,7 +70,7 @@ def cancel_all(update, bot): count += 1 dlDetails._listener.onDownloadError("Download Manually Cancelled By user.") delete_all_messages() - sendMessage(f'Cancelled {count} downloads!', update, bot) + sendMessage(f'Cancelled {count} downloads!', context.bot,update) cancel_mirror_handler = CommandHandler(BotCommands.CancelMirror, cancel_mirror, diff --git a/bot/modules/clone.py b/bot/modules/clone.py index fa33233..407e353 100644 --- a/bot/modules/clone.py +++ b/bot/modules/clone.py @@ -7,15 +7,15 @@ @run_async -def cloneNode(bot,update): +def cloneNode(update,context): args = update.message.text.split(" ",maxsplit=1) if len(args) > 1: link = args[1] - msg = sendMessage(f"Cloning: {link}",bot,update) + msg = sendMessage(f"Cloning: {link}",context.bot,update) gd = GoogleDriveHelper() result = gd.clone(link) - deleteMessage(bot,msg) - sendMessage(result,bot,update) + deleteMessage(context.bot,msg) + sendMessage(result,context.bot,update) else: sendMessage("Provide G-Drive Shareable Link to Clone.",bot,update) diff --git a/bot/modules/list.py b/bot/modules/list.py index bef678c..24a641c 100644 --- a/bot/modules/list.py +++ b/bot/modules/list.py @@ -7,18 +7,18 @@ from bot.helper.telegram_helper.bot_commands import BotCommands @run_async -def list_drive(bot,update): +def list_drive(update,context): message = update.message.text search = message.split(' ',maxsplit=1)[1] LOGGER.info(f"Searching: {search}") gdrive = GoogleDriveHelper(None) msg = gdrive.drive_list(search) if msg: - reply_message = sendMessage(msg, bot, update) + reply_message = sendMessage(msg, context.bot, update) else: - reply_message = sendMessage('No result found', bot, update) + reply_message = sendMessage('No result found', context.bot, update) - threading.Thread(target=auto_delete_message, args=(bot, update.message, reply_message)).start() + threading.Thread(target=auto_delete_message, args=(context.bot, update.message, reply_message)).start() list_handler = CommandHandler(BotCommands.ListCommand, list_drive,filters=CustomFilters.authorized_chat | CustomFilters.authorized_user) diff --git a/bot/modules/mirror.py b/bot/modules/mirror.py index 94ea9ba..9a97645 100644 --- a/bot/modules/mirror.py +++ b/bot/modules/mirror.py @@ -167,13 +167,13 @@ def _mirror(bot, update, isTar=False): @run_async -def mirror(bot, update): - _mirror(bot, update) +def mirror(update, context): + _mirror(context.bot, update) @run_async -def tar_mirror(update, bot): - _mirror(update, bot, True) +def tar_mirror(update, context): + _mirror(context.bot, update, True) mirror_handler = CommandHandler(BotCommands.MirrorCommand, mirror, diff --git a/bot/modules/mirror_status.py b/bot/modules/mirror_status.py index bc2a913..e4ae324 100644 --- a/bot/modules/mirror_status.py +++ b/bot/modules/mirror_status.py @@ -9,11 +9,11 @@ import threading @run_async -def mirror_status(bot,update): +def mirror_status(update,context): message = get_readable_message() if len(message) == 0: message = "No active downloads" - reply_message = sendMessage(message, bot, update) + reply_message = sendMessage(message, context.bot, update) threading.Thread(target=auto_delete_message, args=(bot, update.message, reply_message)).start() return index = update.effective_chat.id @@ -21,8 +21,8 @@ def mirror_status(bot,update): if index in status_reply_dict.keys(): deleteMessage(bot, status_reply_dict[index]) del status_reply_dict[index] - sendStatusMessage(update,bot) - deleteMessage(bot,update.message) + sendStatusMessage(update,context.bot) + deleteMessage(context.bot,update.message) mirror_status_handler = CommandHandler(BotCommands.StatusCommand, mirror_status, diff --git a/requirements.txt b/requirements.txt index 151622b..c5ec760 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,4 @@ -python-telegram-bot==12.2.0 +python-telegram-bot==12.6.1 google-api-python-client>=1.7.11,<1.7.20 google-auth-httplib2>=0.0.3,<0.1.0 google-auth-oauthlib>=0.4.1,<0.10.0 From 63e079b72a93a30b582837bdc6be285dc6eab7f1 Mon Sep 17 00:00:00 2001 From: lzzy12 Date: Tue, 14 Apr 2020 23:20:52 +0530 Subject: [PATCH 37/37] Separate out youtube-dl command - /watch [youtube-dl supported link] Signed-off-by: lzzy12 --- bot/__main__.py | 23 ++++++--- .../youtube_dl_download_helper.py | 38 +++++++++++++-- bot/helper/telegram_helper/bot_commands.py | 3 +- bot/modules/mirror.py | 23 ++++----- bot/modules/watch.py | 47 +++++++++++++++++++ 5 files changed, 110 insertions(+), 24 deletions(-) create mode 100644 bot/modules/watch.py diff --git a/bot/__main__.py b/bot/__main__.py index 0c5c2a4..91db40c 100644 --- a/bot/__main__.py +++ b/bot/__main__.py @@ -8,7 +8,8 @@ import shutil from .helper.telegram_helper.filters import CustomFilters from bot.helper.telegram_helper.bot_commands import BotCommands -from .modules import authorize, list, cancel_mirror, mirror_status, mirror, clone +from .modules import authorize, list, cancel_mirror, mirror_status, mirror, clone, watch + @run_async def stats(update,context): @@ -19,10 +20,9 @@ def stats(update,context): free = get_readable_file_size(free) stats = f'Bot Uptime: {currentTime}\n' \ f'Total disk space: {total}\n' \ - f'Used: {used}\n' \ - f'Free: {free}' - sendMessage(stats, context.bot, update) - + f'Used: {used}\n' \ + f'Free: {free}' + sendMessage(stats, bot, update) @run_async @@ -45,7 +45,12 @@ def log(update,context): @run_async -def bot_help(bot,update): +def log(update,context): + sendLogFile(context.bot, update) + + +@run_async +def bot_help(update, context): help_string = f''' /{BotCommands.HelpCommand}: To get this message @@ -53,6 +58,10 @@ def bot_help(bot,update): /{BotCommands.TarMirrorCommand} [download_url][magnet_link]: start mirroring and upload the archived (.tar) version of the download +/{BotCommands.WatchCommand} [youtube-dl supported link]: Mirror through youtube-dl + +/{BotCommands.TarWatchCommand} [youtube-dl supported link]: Mirror through youtube-dl and tar before uploading + /{BotCommands.CancelMirror} : Reply to the message by which the download was initiated and that download will be cancelled /{BotCommands.StatusCommand}: Shows a status of all the downloads @@ -78,7 +87,7 @@ def main(): help_handler = CommandHandler(BotCommands.HelpCommand, bot_help, filters=CustomFilters.authorized_chat | CustomFilters.authorized_user) stats_handler = CommandHandler(BotCommands.StatsCommand, - stats, filters=CustomFilters.authorized_chat | CustomFilters.authorized_user) + stats, filters=CustomFilters.authorized_chat | CustomFilters.authorized_user) log_handler = CommandHandler(BotCommands.LogCommand, log, filters=CustomFilters.owner_filter) dispatcher.add_handler(start_handler) dispatcher.add_handler(ping_handler) diff --git a/bot/helper/mirror_utils/download_utils/youtube_dl_download_helper.py b/bot/helper/mirror_utils/download_utils/youtube_dl_download_helper.py index e7dfcc9..66e87c1 100644 --- a/bot/helper/mirror_utils/download_utils/youtube_dl_download_helper.py +++ b/bot/helper/mirror_utils/download_utils/youtube_dl_download_helper.py @@ -1,9 +1,32 @@ from .download_helper import DownloadHelper import time -from youtube_dl import YoutubeDL +from youtube_dl import YoutubeDL, DownloadError import threading from bot import LOGGER, download_dict_lock, download_dict, DOWNLOAD_DIR from ..status_utils.youtube_dl_download_status import YoutubeDLDownloadStatus +import logging +import re + +LOGGER = logging.getLogger(__name__) + + +class MyLogger: + def __init__(self, obj): + self.obj = obj + + def debug(self, msg): + LOGGER.debug(msg) + # Hack to fix changing changing extension + match = re.search(r'.ffmpeg..Merging formats into..(.*?).$', msg) + if match and not self.obj.is_playlist: + self.obj.name = match.group(1) + + def warning(self, msg): + LOGGER.warning(msg) + + def error(self, msg): + LOGGER.error(msg) + #self.obj.onDownloadError(msg) class YoutubeDLHelper(DownloadHelper): @@ -16,6 +39,8 @@ def __init__(self, listener): self.opts = { 'format': 'bestaudio/best', 'progress_hooks': [self.__onDownloadProgress], + 'usenetrc': True, + 'format': "best" } self.ydl = YoutubeDL(self.opts) self.__download_speed = 0 @@ -64,7 +89,7 @@ def __onDownloadStart(self): def __onDownloadComplete(self): self.__listener.onDownloadComplete() - def __onDownloadError(self, error): + def onDownloadError(self, error): self.__listener.onDownloadError(error) def extractMetaData(self, link): @@ -90,14 +115,21 @@ def extractMetaData(self, link): def __download(self, link): try: self.ydl.download([link], ) + with YoutubeDL(self.opts) as ydl: + try: + ydl.download([link]) + except DownloadError as e: + self.onDownloadError(str(e)) + return self.__onDownloadComplete() except ValueError: LOGGER.info("Download Cancelled by User!") - self.__onDownloadError("Download Cancelled by User!") + self.onDownloadError("Download Cancelled by User!") def add_download(self, link, path): LOGGER.info(f"Downloading with YT-DL: {link}") self.__gid = f"{self.vid_id}{self.__listener.uid}" + self.opts['logger'] = MyLogger(self) if not self.is_playlist: self.opts['outtmpl'] = f"{path}/%(title)s.%(ext)s" else: diff --git a/bot/helper/telegram_helper/bot_commands.py b/bot/helper/telegram_helper/bot_commands.py index d58c29d..d4e8b64 100644 --- a/bot/helper/telegram_helper/bot_commands.py +++ b/bot/helper/telegram_helper/bot_commands.py @@ -14,6 +14,7 @@ def __init__(self): self.HelpCommand = 'help' self.LogCommand = 'log' self.CloneCommand = "clone" - + self.WatchCommand = 'watch' + self.TarWatchCommand = 'tarwatch' BotCommands = _BotCommands() diff --git a/bot/modules/mirror.py b/bot/modules/mirror.py index 9a97645..0705677 100644 --- a/bot/modules/mirror.py +++ b/bot/modules/mirror.py @@ -30,9 +30,12 @@ def onDownloadProgress(self): pass def clean(self): - Interval[0].cancel() - del Interval[0] - delete_all_messages() + try: + Interval[0].cancel() + del Interval[0] + delete_all_messages() + except IndexError: + pass def onDownloadComplete(self): with download_dict_lock: @@ -65,6 +68,8 @@ def onDownloadComplete(self): drive.upload(up_name) def onDownloadError(self, error): + error = error.replace('<', ' ') + error = error.replace('>', ' ') LOGGER.info(self.update.effective_chat.id) with download_dict_lock: try: @@ -151,16 +156,8 @@ def _mirror(bot, update, isTar=False): sendMessage('No download source provided', bot, update) return listener = MirrorListener(bot, update, isTar, tag) - ydl = YoutubeDLHelper(listener) - try: - sup_link = ydl.extractMetaData(link) - except Exception as e: - sup_link = None - if sup_link: - ydl.add_download(link, f'{DOWNLOAD_DIR}{listener.uid}') - else: - aria = aria2_download.AriaDownloadHelper(listener) - aria.add_download(link, f'{DOWNLOAD_DIR}/{listener.uid}/') + aria = aria2_download.AriaDownloadHelper(listener) + aria.add_download(link, f'{DOWNLOAD_DIR}/{listener.uid}/') sendStatusMessage(update, bot) if len(Interval) == 0: Interval.append(setInterval(DOWNLOAD_STATUS_UPDATE_INTERVAL, update_all_messages)) diff --git a/bot/modules/watch.py b/bot/modules/watch.py new file mode 100644 index 0000000..80148b8 --- /dev/null +++ b/bot/modules/watch.py @@ -0,0 +1,47 @@ +from telegram.ext import CommandHandler, run_async +from telegram import Bot, Update +from bot import Interval, INDEX_URL, DOWNLOAD_DIR, DOWNLOAD_STATUS_UPDATE_INTERVAL, dispatcher, LOGGER +from bot.helper.ext_utils.bot_utils import setInterval +from bot.helper.telegram_helper.message_utils import update_all_messages, sendMessage +from .mirror import MirrorListener +from bot.helper.mirror_utils.download_utils.youtube_dl_download_helper import YoutubeDLHelper +from bot.helper.telegram_helper.bot_commands import BotCommands +from bot.helper.telegram_helper.filters import CustomFilters + + +def _watch(bot: Bot, update: Update, args: list, isTar=False): + try: + link = args[0] + except IndexError: + sendMessage('/watch [yt_dl supported link] to mirror with youtube_dl', bot, update) + return + reply_to = update.message.reply_to_message + if reply_to is not None: + tag = reply_to.from_user.username + else: + tag = None + + listener = MirrorListener(bot, update, isTar, tag) + ydl = YoutubeDLHelper(listener) + ydl.add_download(link, f'{DOWNLOAD_DIR}{listener.uid}') + if len(Interval) == 0: + Interval.append(setInterval(DOWNLOAD_STATUS_UPDATE_INTERVAL, update_all_messages)) + + +@run_async +def watchTar(bot: Bot, update: Update, args: list): + _watch(bot, update, args, True) + + +def watch(bot: Bot, update: Update, args: list): + _watch(bot, update, args) + + +mirror_handler = CommandHandler(BotCommands.WatchCommand, watch, + pass_args=True, + filters=CustomFilters.authorized_chat | CustomFilters.authorized_user) +tar_mirror_handler = CommandHandler(BotCommands.TarWatchCommand, watchTar, + pass_args=True, + filters=CustomFilters.authorized_chat | CustomFilters.authorized_user) +dispatcher.add_handler(mirror_handler) +dispatcher.add_handler(tar_mirror_handler)