Allow for HTTP redirection after successful authentication.

diff --git a/gdrive_knife.py b/gdrive_knife.py
index 94945b281ecbfd30f3df31ecf46b9211c5ee0d7b..b4b99129200d5cc29ffa94ddb4ed3a459511c42c 100644 (file)
--- a/gdrive_knife.py
+++ b/gdrive_knife.py
@@ -91,8 +91,9 @@ def decrypt_chunks_in_place(encryption_key, path):
 
 class auth_server():
 
-    def __init__(self, flow):
+    def __init__(self, flow, redirect_url):
         self.flow = flow
+        self.url = redirect_url
 
     def handle_one_on(self, host, port):
         print(f'Polling on {host}:{port}...')
@@ -104,7 +105,10 @@ class auth_server():
         parameters = urllib.parse.parse_qs(environ['QUERY_STRING'])
         self.flow.fetch_token(code=parameters['code'][0])
 
-        start_response('200 OK', [('Content-type', 'text/html')])
+        if self.url:
+            start_response('308 Permanent Redirect', [('Location', self.url)])
+        else:
+            start_response('200 OK', [('Content-type', 'text/html')])
         return [b'Obtainted new credentials.']
 
 def auth(args):
@@ -125,9 +129,10 @@ def auth(args):
         else:
             print(f'{args.token} does not exist. Obtaining a new token.')
             flow = Flow.from_client_secrets_file(args.credentials, scopes=SCOPES,
-                redirect_uri=args.redirect_url)
+                redirect_uri=args.on_token)
             # Run the server on localhost because wsgi does not use HTTPS.
-            auth_server(flow).handle_one_on('localhost', args.port)
+            auth_server(flow, args.on_success).handle_one_on('localhost',
+                args.port)
             creds = flow.credentials
         print(f'Writing new token to {args.token}.')
         with open(args.token, 'w') as token:
@@ -235,7 +240,9 @@ if __name__ == '__main__':
         type=lambda x : file_path(parser, x), help='File with credentials.')
     auth_parser.add_argument('-p', '--port', type=int, default=8080, help='Port '
         'for the authentication server.')
-    auth_parser.add_argument('-r', '--redirect', default='http://localhost:8080',
+    auth_parser.add_argument('--on_success', help='URL to redirect to after '
+        'successful authentication.')
+    auth_parser.add_argument('--on_token', default='http://localhost:8080',
         help='URL to redirect to after successful authentication.')
     auth_parser.add_argument('-t', '--token', default='token.json',
         required=True, help='File with the authentication token (created if '