tools/lib/auth.py: support github/apple login (#22766)

* tools/lib/auth.py: support github/apple login

* print some info about logged in user

* add docstring

tools/lib/auth.py

@@ -1,20 +1,47 @@
 #!/usr/bin/env python3
+"""
+Usage::
+
+  usage: auth.py [-h] [{google,apple,github,jwt}] [jwt]
+
+  Login to your comma account
+
+  positional arguments:
+    {google,apple,github,jwt}
+    jwt
+
+  optional arguments:
+    -h, --help            show this help message and exit
+
+
+Examples::
+
+  ./auth.py  # Log in with google account
+  ./auth.py github  # Log in with GitHub Account
+  ./auth.py jwt ey......hw  # Log in with a JWT from https://jwt.comma.ai, for use in CI
+"""
+
+import argparse
 import sys
+import pprint
 import webbrowser
-from http.server import HTTPServer, BaseHTTPRequestHandler
-from urllib.parse import urlencode, parse_qs
-from tools.lib.api import CommaApi, APIError
-from tools.lib.auth_config import set_token
-from typing import Dict, Any
+from http.server import BaseHTTPRequestHandler, HTTPServer
+from typing import Any, Dict
+from urllib.parse import parse_qs, urlencode
+
+from tools.lib.api import APIError, CommaApi, UnauthorizedError
+from tools.lib.auth_config import set_token, get_token
 
 PORT = 3000
 
+
 class ClientRedirectServer(HTTPServer):
   query_params: Dict[str, Any] = {}
 
+
 class ClientRedirectHandler(BaseHTTPRequestHandler):
   def do_GET(self):
-    if not self.path.startswith('/auth/g/redirect'):
+    if not self.path.startswith('/auth'):
       self.send_response(204)
       return
 
@@ -30,21 +57,48 @@ class ClientRedirectHandler(BaseHTTPRequestHandler):
   def log_message(self, format, *args):  # pylint: disable=redefined-builtin
     pass  # this prevent http server from dumping messages to stdout
 
-def auth_redirect_link():
-  redirect_uri = f'http://localhost:{PORT}/auth/g/redirect'
+
+def auth_redirect_link(method):
+  provider_id = {
+    'google': 'g',
+    'apple': 'a',
+    'github': 'h',
+  }[method]
+
   params = {
+    'redirect_uri': f"https://api.comma.ai/v2/auth/{provider_id}/redirect/",
+    'state': f'service,localhost:{PORT}',
+  }
+
+  if method == 'google':
+    params.update({
       'type': 'web_server',
       'client_id': '45471411055-ornt4svd2miog6dnopve7qtmh5mnu6id.apps.googleusercontent.com',
-    'redirect_uri': redirect_uri,
       'response_type': 'code',
       'scope': 'https://www.googleapis.com/auth/userinfo.email',
       'prompt': 'select_account',
-  }
+    })
+    return 'https://accounts.google.com/o/oauth2/auth?' + urlencode(params)
+  elif method == 'github':
+    params.update({
+      'client_id': '28c4ecb54bb7272cb5a4',
+      'scope': 'read:user',
+    })
+    return 'https://github.com/login/oauth/authorize?' + urlencode(params)
+  elif method == 'apple':
+      params.update({
+        'client_id': 'ai.comma.login',
+        'response_type': 'code',
+        'response_mode': 'form_post',
+        'scope': 'name email',
+      })
+      return 'https://appleid.apple.com/auth/authorize?' + urlencode(params)
+  else:
+    raise NotImplementedError(f"no redirect implemented for method {method}")
 
-  return (redirect_uri, 'https://accounts.google.com/o/oauth2/auth?' + urlencode(params))
 
-def login():
-  redirect_uri, oauth_uri = auth_redirect_link()
+def login(method):
+  oauth_uri = auth_redirect_link(method)
 
   web_server = ClientRedirectServer(('localhost', PORT), ClientRedirectHandler)
   print(f'To sign in, use your browser and navigate to {oauth_uri}')
@@ -53,7 +107,6 @@ def login():
   while True:
     web_server.handle_request()
     if 'code' in web_server.query_params:
-      code = web_server.query_params['code']
       break
     elif 'error' in web_server.query_params:
       print('Authentication Error: "%s". Description: "%s" ' % (
@@ -62,11 +115,31 @@ def login():
       break
 
   try:
-    auth_resp = CommaApi().post('v2/auth/', data={'code': code, 'redirect_uri': redirect_uri})
+    auth_resp = CommaApi().post('v2/auth/', data={'code': web_server.query_params['code'], 'provider': web_server.query_params['provider']})
     set_token(auth_resp['access_token'])
-    print('Authenticated')
   except APIError as e:
     print(f'Authentication Error: {e}', file=sys.stderr)
 
+
 if __name__ == '__main__':
-  login()
+  parser = argparse.ArgumentParser(description='Login to your comma account')
+  parser.add_argument('method', default='google', const='google', nargs='?', choices=['google', 'apple', 'github', 'jwt'])
+  parser.add_argument('jwt', nargs='?')
+
+  args = parser.parse_args()
+  if args.method == 'jwt':
+    if args.jwt is None:
+      print("method JWT selected, but no JWT was provided")
+      exit(1)
+
+    set_token(args.jwt)
+  else:
+    login(args.method)
+
+  try:
+    me = CommaApi(token=get_token()).get('/v1/me')
+    print("Authenticated!")
+    pprint.pprint(me)
+  except UnauthorizedError:
+    print("Got invalid JWT")
+    exit(1)