Add support for service account

Adding support service accounts

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

Added scripts and docs for generating service accounts

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

gen_sa_accounts: Save credentials with indexed file name

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

gdriveTools: Avoid using oauth2 library for service accounts

oauth2 library is deprecated

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

Author: lzzy12

.gitignore

@@ -10,3 +10,4 @@ data*
 *.pickle
 authorized_chats.txt
 log.txt
+accounts/*

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`[email protected]`.
+
+3. Run `python3 add_to_google_group.py -g [email protected]`
+
+4. Now, add Google Groups (**Step 2**) to manager your service accounts, add the group address `[email protected]` or `[email protected]` to the Team drive or folder

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  [email protected]
+
+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')
+# [email protected]
+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))

bot/__init__.py

@@ -90,6 +90,16 @@ def getConfig(name: str):
         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
+dispatcher = updater.dispatcher

bot/helper/mirror_utils/upload_utils/gdriveTools.py

@@ -4,33 +4,66 @@
 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 @@ 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 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 @@ 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))
     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,34 +151,42 @@ 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(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,
                                      resumable=True,
                                      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 @@ def cloneFolder(self,name,local_path,folder_id,parent_id):
     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,40 +338,20 @@ 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.
         query = f"'{parent_id}' in parents and (name contains '{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

config_sample.env

@@ -12,4 +12,5 @@ IS_TEAM_DRIVE = ""
 INDEX_URL = ""
 USER_SESSION_STRING = ""
 TELEGRAM_API = 
-TELEGRAM_HASH = ""
+TELEGRAM_HASH = ""
+USE_SERVICE_ACCOUNTS = ""