Permit remote redirect URL to allow for remote authentication.

diff --git a/gdrive_knife.py b/gdrive_knife.py
index af0bda7006d81b5afaead597c4dda8d202678946..91782101fbca2a570994e32a16e5b03ef361040e 100644 (file)
--- a/gdrive_knife.py
+++ b/gdrive_knife.py
@@ -6,6 +6,7 @@ import io
 import os.path
 import shutil
 import sys
+import urllib.parse
 import uuid
 import zipfile
 
@@ -13,9 +14,10 @@ from apiclient.http import MediaFileUpload
 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 google_auth_oauthlib.flow import Flow
 from googleapiclient.discovery import build
 from googleapiclient.http import MediaIoBaseDownload
+from wsgiref.simple_server import make_server
 
 SCOPES = ['https://www.googleapis.com/auth/drive.metadata',
           'https://www.googleapis.com/auth/drive',
@@ -87,20 +89,46 @@ def decrypt_chunks_in_place(encryption_key, path):
             output_file.write(decrypted_bytes)
     os.rename(tmp_path, path)
 
+class auth_server():
+
+    def __init__(self, flow):
+        self.flow = flow
+
+    def handle_one_on(self, host, port):
+        print(f'Polling on {host}:{port}...')
+        with make_server(host, port, self) as httpd:
+            print(f'Authentication URL: {self.flow.authorization_url()[0]}')
+            httpd.handle_request()
+
+    def __call__(self, environ, start_response):
+        parameters = urllib.parse.parse_qs(environ['QUERY_STRING'])
+        self.flow.fetch_token(code=parameters['code'][0])
+
+        start_response('200 OK', [('Content-type', 'text/html')])
+        return [b'Obtainted new credentials.']
+
 def auth(args):
     creds = None
     if os.path.exists(args.token):
         print(f'{args.token} exists. Trying to authenticate with it.')
-        creds = Credentials.from_authorized_user_file(args.token, SCOPES)
+        try:
+            creds = Credentials.from_authorized_user_file(args.token, SCOPES)
+        except ValueError:
+            os.remove(args.token)
+            print(f'{args.token} is malformed. Reset permissions at '
+                   'https://myaccount.google.com/permissions and retry.')
+            sys.exit(1)
     if not creds or not creds.valid:
         if creds and creds.expired and creds.refresh_token:
             print(f'{args.token} has expired. Refreshing.')
             creds.refresh(Request())
         else:
             print(f'{args.token} does not exist. Obtaining a new token.')
-            flow = InstalledAppFlow.from_client_secrets_file(args.credentials,
-                SCOPES)
-            creds = flow.run_local_server(port=0)
+            flow = Flow.from_client_secrets_file(args.credentials, scopes=SCOPES,
+                redirect_uri=args.redirect_url)
+            # Run the server on localhost because wsgi does not use HTTPS.
+            auth_server(flow).handle_one_on('localhost', args.port)
+            creds = flow.credentials
         print(f'Writing new token to {args.token}.')
         with open(args.token, 'w') as token:
             token.write(creds.to_json())
@@ -205,6 +233,11 @@ if __name__ == '__main__':
         'token.')
     auth_parser.add_argument('-c', dest='credentials', required=True,
         type=lambda x : file_path(parser, x), help='File with credentials.')
+    auth_parser.add_argument('-p', dest='port', type=int, default=8080,
+        help='Port on which the authentication server listens.')
+    auth_parser.add_argument('-r', dest='redirect_url',
+        default='http://localhost:8080', help='URL to redirect to after '
+        'successful after successful authentication.')
     auth_parser.add_argument('-t', dest='token', default='token.json',
         required=True, help='Path where the authentication token should be '
         'saved.')