From: Jakub Czajka <jakub@ekhem.eu.org>
Date: Tue, 12 Sep 2023 18:09:39 +0000 (+0200)
Subject: Add a subprogram for downloading files from the drive.
X-Git-Url: https://git.ekhem.eu.org/?a=commitdiff_plain;h=c2ff6b1b1934f686335fad07b9f71285e18b4cac;p=gdrive_knife.git

Add a subprogram for downloading files from the drive.
---

diff --git a/gdrive_knife.py b/gdrive_knife.py
index 9543ab9..0d72528 100644
--- a/gdrive_knife.py
+++ b/gdrive_knife.py
@@ -2,13 +2,20 @@
 # License: GPL-3.0 or later.
 
 import argparse
+import io
 import os.path
+import shutil
 import sys
+import tempfile
+import uuid
+import zipfile
 
+from cryptography.fernet import Fernet
 from google.auth.transport.requests import Request
 from google.oauth2.credentials import Credentials
 from google_auth_oauthlib.flow import InstalledAppFlow
 from googleapiclient.discovery import build
+from googleapiclient.http import MediaIoBaseDownload
 
 SCOPES = ['https://www.googleapis.com/auth/drive.metadata',
           'https://www.googleapis.com/auth/drive',
@@ -19,6 +26,12 @@ def file_path(parser, path):
         return parser.error(f'{path} does not exist!')
     return path
 
+def fernet_key(parser, path):
+    if not os.path.isfile(path):
+        return parser.error(f'{path} does not exist!')
+    with open(path, 'rb') as f:
+        return Fernet(f.read())
+
 def get_drive_client(authentication_token):
     creds = Credentials.from_authorized_user_file(authentication_token, SCOPES)
     if not creds or not creds.valid:
@@ -28,6 +41,20 @@ def get_drive_client(authentication_token):
 
     return build('drive', 'v3', credentials=creds)
 
+def get_path_on_drive(file_path):
+    if os.path.isabs(file_path):
+        return tempfile.gettempdir() + file_path
+    return tempfile.gettempdir() + '/' + file_path
+
+def get_file_id(drive, file_path):
+    path_on_drive = get_path_on_drive(file_path)
+    query_result = drive.files().list(q=f"name='{path_on_drive}'",
+        fields='files(id)').execute()
+    maybe_id = query_result.get('files', [])
+    if not maybe_id:
+        return None
+    return maybe_id[0]['id']
+
 def auth(args):
     creds = None
     if os.path.exists(args.token):
@@ -56,6 +83,39 @@ def list(args):
     for file_on_drive in files_on_drive['files']:
         print(file_on_drive['originalFilename'])
 
+def download(args):
+    drive = get_drive_client(args.token)
+
+    maybe_id = get_file_id(drive, args.path)
+    if not maybe_id:
+        print(f'File {args.path} not found.')
+        sys.exit(1)
+
+    request = drive.files().get_media(fileId=maybe_id, acknowledgeAbuse=True)
+    file = io.BytesIO()
+    downloader = MediaIoBaseDownload(file, request)
+    done = False
+    while done is False:
+        status, done = downloader.next_chunk()
+        print(F'Download {int(status.progress() * 100)}.')
+
+    encrypted_file = file.getvalue()
+    token = args.key.decrypt(encrypted_file)
+    print(f'{args.path} decrypted.')
+
+    path_in_tmp = tempfile.gettempdir() + '/' + str(uuid.uuid4())
+    with open(path_in_tmp, 'wb+') as outfile:
+        outfile.write(token)
+    print(f'{args.path} written to {path_in_tmp}.')
+
+    if zipfile.is_zipfile(path_in_tmp):
+        os.makedirs(args.output, exist_ok=True)
+        shutil.unpack_archive(path_in_tmp, extract_dir=args.output, format='zip')
+        print(f'Unarchived to {args.output}.')
+    else:
+        shutil.move(path_in_tmp, args.output)
+        print(f'Moved to {args.output}.')
+
 if __name__ == '__main__':
     parser = argparse.ArgumentParser(prog='gdrive_knife', description='Swiss '
         'army knife for working with Google Drive.')
@@ -77,6 +137,18 @@ if __name__ == '__main__':
         'authentication token.')
     list_parser.set_defaults(func=list)
 
+    download_parser = subparsers.add_parser('download', help='Download a file.')
+    download_parser.add_argument('path', help='Path on the drive to download.')
+    download_parser.add_argument('output', help='Path where the file should be '
+        'saved.')
+    download_parser.add_argument('-k', dest='key', required=True,
+        type=lambda x : fernet_key(parser, x), help='File with the '
+        'decryption key.')
+    download_parser.add_argument('-t', dest='token', required=True,
+        type=lambda x : file_path(parser, x), help='File with the '
+        'authentication token.')
+    download_parser.set_defaults(func=download)
+
     args = parser.parse_args()
 
     args.func(args)
diff --git a/requirements.txt b/requirements.txt
index 67480db..adc302e 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,3 +1,4 @@
+cryptography==40.0.2
 google-api-python-client==2.84.0
 google-auth-httplib2==0.1.0
 google-auth-oauthlib==1.0.0