From 7ad7d9388adec4584690ca9898dcfd4e160428b4 Mon Sep 17 00:00:00 2001 From: lzzy12 Date: Tue, 21 Jan 2020 03:49:15 -0800 Subject: [PATCH] Add support for service account Adding support service accounts Signed-off-by: lzzy12 Added scripts and docs for generating service accounts Signed-off-by: lzzy12 gen_sa_accounts: Save credentials with indexed file name Signed-off-by: lzzy12 gdriveTools: Avoid using oauth2 library for service accounts oauth2 library is deprecated Signed-off-by: lzzy12 --- .gitignore | 1 + README.md | 31 ++ add_to_google_group.py | 84 ++++ bot/__init__.py | 12 +- .../mirror_utils/upload_utils/gdriveTools.py | 138 ++++--- config_sample.env | 3 +- gen_sa_accounts.py | 365 ++++++++++++++++++ 7 files changed, 582 insertions(+), 52 deletions(-) create mode 100644 add_to_google_group.py create mode 100644 gen_sa_accounts.py 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/README.md b/README.md index be98fa8..f0db9b9 100644 --- a/README.md +++ b/README.md @@ -87,3 +87,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/bot/__init__.py b/bot/__init__.py index 7190675..1fd6ba9 100644 --- a/bot/__init__.py +++ b/bot/__init__.py @@ -90,6 +90,16 @@ try: 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 \ No newline at end of file +dispatcher = updater.dispatcher diff --git a/bot/helper/mirror_utils/upload_utils/gdriveTools.py b/bot/helper/mirror_utils/upload_utils/gdriveTools.py index 18eeea3..62bd6ef 100644 --- a/bot/helper/mirror_utils/upload_utils/gdriveTools.py +++ b/bot/helper/mirror_utils/upload_utils/gdriveTools.py @@ -4,33 +4,66 @@ import urllib.parse as urlparse from urllib.parse import parse_qs import requests + 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 tenacity import * -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, \ + 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 = 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" 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.__G_DRIVE_DIR_BASE_DOWNLOAD_URL = "https://drive.google.com/drive/folders/{}" self.__listener = listener - self.__service = self.authorize() self._file_uploaded_bytes = 0 self.uploaded_bytes = 0 self.start_time = 0 @@ -74,6 +107,21 @@ class GoogleDriveHelper: 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 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): @@ -83,11 +131,13 @@ class GoogleDriveHelper: '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)) 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, @@ -101,13 +151,13 @@ class GoogleDriveHelper: 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(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 media_body = MediaFileUpload(file_path, mimetype=mime_type, @@ -115,20 +165,28 @@ class GoogleDriveHelper: 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): @@ -249,11 +307,11 @@ class GoogleDriveHelper: 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) @@ -280,26 +338,6 @@ class GoogleDriveHelper: 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. @@ -307,13 +345,13 @@ class GoogleDriveHelper: 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 diff --git a/config_sample.env b/config_sample.env index b0f5156..1ef1dc7 100644 --- a/config_sample.env +++ b/config_sample.env @@ -12,4 +12,5 @@ IS_TEAM_DRIVE = "" INDEX_URL = "" USER_SESSION_STRING = "" TELEGRAM_API = -TELEGRAM_HASH = "" \ No newline at end of file +TELEGRAM_HASH = "" +USE_SERVICE_ACCOUNTS = "" diff --git a/gen_sa_accounts.py b/gen_sa_accounts.py new file mode 100644 index 0000000..0fd1c24 --- /dev/null +++ b/gen_sa_accounts.py @@ -0,0 +1,365 @@ +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.')