#!/usr/bin/env python3
#
#
# Copyright 2021 The Android Open Source Project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Uploader tool for sending device keys to the Android Partner Front End.

This tool consumes an input file containing device info, which includes the
device's public key, and batch uploads it to Google. Once uploaded, the device
may use remote key provisioning to request attestation certificates from Google.

This tool is designed to be used with the rkp_factory_extraction_tool.
Therefore, the JSON output from rkp_factory_extraction_tool is the expected
input, with one JSON string per line of input.

This script requires an oauth2 client id to authenticate to the endpoint. See
https://support.google.com/cloud/answer/6158849 for more info on setting up
a client id.
"""

import argparse
import http.server
import json
import os
import sys
import urllib.parse
import urllib.request
import webbrowser

GOOGLE_AUTH_CLIENT_INSTALLED = False
try:
  from google.auth.transport import requests
  from google.oauth2 import service_account

  GOOGLE_AUTH_CLIENT_INSTALLED = True
except ImportError:
  pass

###############################################################################
# Release Notes:
#
# February 15, 2023:
# - Add support for BatchCheck API.
#
# January 26, 2023:
# - Updated error/warning handling.
#
# November 1, 2022:
# - Added support for "--no-update" option.
#
# October 31, 2022:
# - Added support for service account based authentication.
# - Prints error messages even if the server sends an OK HTTP status.
#
# September 15, 2022:
# - The script performs an "upsert" by default. Duplicate keys are updated,
#   and missing keys are inserted. Deprecated "--update" option.
#
# September 7, 2022:
# - Use the v1 uploader API, resulting in:
#   - The build fingerprint is no longer necessary.
#   - Callers must supply a company identifier
# - Upload CSRs in batches of 2000.
###############################################################################

DEFAULT_BASE = 'https://androidpartner.googleapis.com'
UPDATE_PATH = '/v1/devices:batchUpdate'
UPLOAD_PATH = '/v1/devices:batchUpload'
BATCH_CHECK_PATH = '/v1/devices:batchCheck'
MAX_BATCH_SIZE = 2000
TOKEN_CACHE_FILE = os.path.join(
    os.path.expanduser('~'), '.device_info_uploader.token'
)

OAUTH_SCOPE = 'https://www.googleapis.com/auth/androidPartner'
OAUTH_SERVICE_BASE = 'https://accounts.google.com/o/oauth2'
OAUTH_AUTHN_URL = OAUTH_SERVICE_BASE + '/auth'
OAUTH_TOKEN_URL = OAUTH_SERVICE_BASE + '/token'
SERVICE_ACCOUNT = None


class OAuthHTTPRequestHandler(http.server.BaseHTTPRequestHandler):
  """HTTP Handler used to accept the oauth response when the user logs in."""

  def do_GET(self):  # pylint: disable=invalid-name
    """Handles GET, extracting the authorization code in the query params."""
    print(f'GET path: {self.path}')
    parsed_path = urllib.parse.urlparse(self.path)
    params = dict(urllib.parse.parse_qsl(parsed_path.query))
    if 'error' in params:
      error = params['error']
      self.respond(
          400, error, f'Error received from the OAuth server: {error}.'
      )
      sys.exit(-1)
    elif 'code' not in params:
      self.respond(
          400,
          'ERROR',
          (
              'Response from OAuth server is missing the authorization '
              f'code. Full response: "{self.path}"'
          ),
      )
      sys.exit(-1)
    else:
      self.respond(
          200, 'Success!', 'Success! You may close this browser window.'
      )
      self.server.code = params['code']

  def do_POST(self):  # pylint: disable=invalid-name
    print(f'POST path: {self.path}')

  def respond(self, code, title, message):
    """Send a response to the HTTP client.

    Args:
        code: The HTTP status code to send
        title: The page title to display
        message: The message to display to the user on the page
    """
    if code != 200:
      eprint(message)
    self.send_response(code)
    self.send_header('Content-type', 'text/html')
    self.end_headers()
    self.wfile.write(
        (
            '<html>'
            f'  <title>{title}</title>'
            '  <body>'
            f'    <p style="font-size:24px;">{message}</p>'
            '  </body>'
            '</html>'
        ).encode('utf-8')
    )


class LocalOAuthReceiver(http.server.HTTPServer):
  """HTTP server that will wait for an OAuth authorization code."""

  def __init__(self):
    super(LocalOAuthReceiver, self).__init__(
        ('127.0.0.1', 0), OAuthHTTPRequestHandler
    )
    self.code = None

  def port(self):
    return self.socket.getsockname()[1]

  def wait_for_code(self):
    print('Waiting for a response from the Google OAuth service.')
    print('If you receive an error in your browser, interrupt this script.')
    self.handle_request()
    return self.code


def eprint(message):
  print(message, file=sys.stderr)


def die(message):
  eprint('ERROR:')
  eprint(message)
  sys.exit(-1)


def parse_args():
  """Parse and return the command line args.

  Returns:
    An argparse.Namespace object populated with the arguments.
  """
  parser = argparse.ArgumentParser(description='Upload device info')
  parser.add_argument(
      '--json-csr',
      nargs='+',
      required=True,
      help=(
          'list of files containing JSON output from'
          ' rkp_factory_extraction_tool'
      ),
  )
  parser.add_argument(
      '--company-id',
      required=True,
      help=(
          'APM company identifier. This value can be found on the Android\n'
          "Partner Portal under Builds -> Metadata. If you can't locate\n"
          'this value, please contact your TAM.\n'
      ),
  )

  credential_arg_group = parser.add_mutually_exclusive_group(required=True)
  credential_arg_group.add_argument(
      '--credentials',
      required=False,
      help='JSON credentials file for web browser OAuth2 authentication',
  )
  credential_arg_group.add_argument(
      '--credentials-keyfile',
      required=False,
      help='JSON service account key file for headless authentication',
  )

  parser.add_argument(
      '--endpoint', default=DEFAULT_BASE, help='destination server URL'
  )
  parser.add_argument(
      '--cache-token',
      action='store_true',
      help=(
          'Use a locally cached refresh token (only applies when using'
          ' --credentials)'
      ),
  )
  parser.add_argument(
      '--test-build',
      action='store_true',
      help=(
          'Use the test build fingerprint e.g. in case of testing an unreleased'
          ' build'
      ),
  )

  # Update and check are exclusive actions.
  exclusive_group = parser.add_mutually_exclusive_group(required=False)
  exclusive_group.add_argument(
      '--update',
      action='store_true',# argparse.BooleanOptionalAction,
      help=(
          'Whether to perform an upsert on the CSRs. Upsert is the default'
          ' behavior.'
      ),
  )
  exclusive_group.add_argument(
      '--check',
      action='store_true',
      required=False,
      help='Perform a batch check on the CSRs.',
  )

  parser.add_argument(
      '--is-test',
      action='store_true',
      help=(
          'Marks the uploaded device as being test only. Used for testing\n'
          'the end-to-end workflow of registering devices, and the signed\n'
          'certificate will not work for attestation.\n'
          'NOTE: "test" status is permanent for a device. A device that has\n'
          "      been registered can't change status between test <-> prod."
      ),
  )

  parser.add_argument(
      '--verbose',
      action='store_true',
      required=False,
      help=(
          'Print request details, and registered device information when'
          ' check-only is true.'
      ),
  )
  args = parser.parse_args()

  return args


def parse_json_csrs(filename, batches):
  """Parse the given file and insert it into batches.

  If the input is not a valid JSON CSR blob, exit the program.

  Args:
    filename: The file that contains a JSON-formatted build and CSR
    batches: Output dict containing a mapping of CSRs to fingerprints
  """
  line_count = 0
  for line in open(filename):
    if line_count % MAX_BATCH_SIZE == 0:
      current_batch = []
      batches.append(current_batch)

    line_count = line_count + 1
    try:
      obj = json.loads(line)
    except json.JSONDecodeError as e:
      die(f'{e.msg} {filename}:{line_count}, char {e.pos}')

    try:
      current_batch.append(obj['csr'])
    except KeyError as e:
      die(f'Invalid object at {filename}:{line_count}, missing {e}')


def format_request_body(args, csrs, upsert):
  """Generate a formatted request buffer for the given build and CSRs."""
  request = {
      'device_payloads': [
          {
              'certificate_payloads': [
                  {'certificate_signing_request': c} for c in csrs
              ]
          }
      ],
      'company_id': args.company_id,
      'is_test': args.is_test,
  }

  if upsert:
    request['allow_missing'] = True

  if args.verbose:
    print('request body:')
    print(request)

  return json.dumps(request).encode('utf-8')


def format_check_request_body(args, csrs):
  """Generate a formatted BatchCheck request buffer for the given build and CSRs."""
  request = {
      'certificate_signing_requests': [c for c in csrs],
      'company_id': args.company_id,
      'return_registered_csrs': args.verbose,
  }

  if args.verbose:
    print('request body:')
    print(request)

  return json.dumps(request).encode('utf-8')


def load_refresh_token():
  if not os.path.exists(TOKEN_CACHE_FILE):
    return None
  with open(TOKEN_CACHE_FILE) as f:
    return f.readline()


def store_refresh_token(refresh_token):
  with open(TOKEN_CACHE_FILE, 'w') as f:
    f.write(refresh_token)


def fetch_access_token(creds, cache_token=False, code=None, redirect_uri=None):
  """Fetch an oauth2 access token.

  If a code is passed, then it is used to get the token. If code
  is None, then look for a persisted refresh token and use that to
  get the access token instead.

  Args:
    creds: The OAuth client credentials, including client secret and id.
    cache_token: If True, then the refresh token is cached on disk so that the
      user does not have to reauthenticate when the script is used again.
    code: The OAuth authorization code, returned by Google's OAuth service.
    redirect_uri: If an authorization code is supplied, then the redirect_uri
      used to fetch the code must be passed here.

  Returns:
    A base64-encode OAuth access token, suitable for including in a request.
  """
  request = urllib.request.Request(OAUTH_TOKEN_URL)
  request.add_header('Content-Type', 'application/x-www-form-urlencoded')
  body = 'client_id=' + creds['client_id']
  body += '&client_secret=' + creds['client_secret']

  if code is not None:
    if redirect_uri is None:
      raise ValueError('"code" was supplied, but "redirect_uri" is None')
    body += '&grant_type=authorization_code'
    body += '&code=' + code
    body += '&redirect_uri=' + redirect_uri
  else:
    refresh_token = load_refresh_token()
    if refresh_token is None:
      return None
    body += '&grant_type=refresh_token'
    body += '&refresh_token=' + refresh_token

  try:
    response = urllib.request.urlopen(request, body.encode('utf-8'))
    parsed_response = json.load(response)
    if cache_token:
      store_refresh_token(parsed_response['refresh_token'])
    return parsed_response['access_token']
  except urllib.error.HTTPError as e:
    # Catch bogus/expired refresh tokens, but bubble up errors when
    # an authorization code is used.
    if code is None:
      return None
    die(f'Failed to receive access token: {e.code} {e.reason}')


def load_and_validate_creds(credfile):
  """Loads the credentials from the given file and validates them.

  Args:
    credfile: the name of the file containing the client credentials

  Returns:
    A map containing the credentials for connecting to the APE backend.
  """
  credmap = json.load(open(credfile))

  not_local_app_creds_error = (
      'ERROR: Invalid credential file.\n'
      '  The given credentials do not appear to be for a locally installed\n'
      '  application. Please navigate to the credentials dashboard and\n'
      '  ensure that the "Type" of your client is "Desktop":\n'
      '  https://console.cloud.google.com/apis/credentials'
  )

  if 'installed' not in credmap:
    die(not_local_app_creds_error)

  creds = credmap['installed']

  expected_keys = set(['client_id', 'client_secret', 'redirect_uris'])
  if not expected_keys.issubset(creds.keys()):
    die(
        'ERROR: Invalid credential file.\n'
        '  The given credentials do not appear to be valid. Please\n'
        '  re-download the client credentials file from the dashboard:\n'
        '  https://console.cloud.google.com/apis/credentials'
    )

  if 'http://localhost' not in creds['redirect_uris']:
    die(not_local_app_creds_error)

  return creds


def authenticate_and_fetch_token(args):
  """Authenticate the user and fetch an access token."""

  def use_service_account():
    """Use a service account to get an access token."""
    global SERVICE_ACCOUNT

    if SERVICE_ACCOUNT is None:
      SERVICE_ACCOUNT = service_account.Credentials.from_service_account_file(
          args.credentials_keyfile,
          scopes=[OAUTH_SCOPE],
      )
      SERVICE_ACCOUNT.refresh(requests.Request())

    return SERVICE_ACCOUNT.token

  def use_oauth2_and_browser():
    """Use OAuth2 credentials to get an access token. Requires a web browser."""
    creds = load_and_validate_creds(args.credentials)

    access_type = 'online'
    if args.cache_token:
      token = fetch_access_token(creds)

      if token is not None:
        return token

      access_type = 'offline'

    httpd = LocalOAuthReceiver()
    redirect_uri = f'http://127.0.0.1:{httpd.port()}'
    url = (
        OAUTH_AUTHN_URL
        + '?response_type=code'
        + '&client_id='
        + creds['client_id']
        + '&redirect_uri='
        + redirect_uri
        + '&scope='
        + OAUTH_SCOPE
        + '&access_type='
        + access_type
        + '&prompt=select_account'
    )
    print('Opening your web browser to authenticate...')
    if not webbrowser.open(url, new=1, autoraise=True):
      print('Error opening the browser. Please open this link in a browser')
      print(f'that is running on this same system:\n  {url}\n')

    code = httpd.wait_for_code()
    return fetch_access_token(creds, args.cache_token, code, redirect_uri)

  if args.credentials_keyfile:
    token = use_service_account()
  elif args.credentials:
    token = use_oauth2_and_browser()
  else:
    raise RuntimeError('Missing credential arg. Should be unreachable')

  return token


def batch_action(args, path, body):
  """Batch action (upload or check existence) for all the CSRs in chunks.

  Args:
    args: The parsed command-line arguments
    csrs: a list of CSRs to be uploaded/checked for the given build
    path: The endpoint url for the specific action
    body: The formatted request body
  """
  request = urllib.request.Request(args.endpoint + path)
  request.add_header('Content-Type', 'application/json')
  request.add_header('X-GFE-SSL', 'yes')
  request.add_header(
      'Authorization', 'Bearer ' + authenticate_and_fetch_token(args)
  )
  try:
    response = urllib.request.urlopen(request, body)
    response_text = response.read().decode('utf-8')
    response_json = json.loads(response_text)

    eprint('Batch API response:' + response_text)

    # It's possible for the API to throw an error but not set an error HTTP code
    if (
        'error' in response_json
        or 'errors' in response_json
        or 'csr_errors' in response_json
    ):
      eprint(f'Errors occurred when uploading csrs: {response.getcode()}')
      eprint(response_text)
      eprint(
          'Refer to '
          + 'https://docs.partner.android.com/gms/building/integrating/att-keys/rkp-errors '
          + 'for detailed explanations of each error code.'
      )
      sys.exit(1)

    if 'warnings' in response_json:
      eprint(f'Warnings occurred when uploading csrs: {response.getcode()}')
      eprint(response_json['warnings'])
  except urllib.error.HTTPError as e:
    eprint(f'Error uploading csrs. {e}')
    for line in e:
      eprint(line.decode('utf-8').rstrip())
    eprint(
        'Refer to '
        + 'https://docs.partner.android.com/gms/building/integrating/att-keys/rkp-errors '
        + 'for detailed explanations of each error code.'
    )
    sys.exit(1)


def upload_batch(args, csrs, upsert):
  """Batch upload all the CSRs in chunks.

  Args:
    args: The parsed command-line arguments
    csrs: a list of CSRs to be uploaded for the given build
    upsert: When True will update CSRs if they have already been uploaded
  """
  body = format_request_body(args, csrs, upsert)
  path = UPDATE_PATH if upsert else UPLOAD_PATH
  batch_action(args, path, body)


def check_batch(args, csrs):
  """Batch check all the CSRs in chunks.

  Args:
    args: The parsed command-line arguments
    csrs: a list of CSRs to be uploaded for the given build
  """
  body = format_check_request_body(args, csrs)
  path = BATCH_CHECK_PATH
  batch_action(args, path, body)


def main():
  args = parse_args()

  if args.credentials_keyfile and not GOOGLE_AUTH_CLIENT_INSTALLED:
    die(
        '--credentials-keyfile argument was specified but you have not '
        "installed Google's auth client library.\n"
        'Run the following command to install it:\n'
        'pip3 install google-auth==2.13.0 requests==2.28'
    )

  check = args.check is True
  upsert = args.update is not False

  batches = []
  for filename in args.json_csr:
    parse_json_csrs(filename, batches)

  batch_count = 1
  for batch in batches:
    if check:
      print(
          '{}/{}: Checking batch containing {} csr(s)'.format(
              batch_count, len(batches), len(batch)
          )
      )
      check_batch(args, batch)
    else:
      print(
          '{}/{}: Uploading batch containing {} csr(s)'.format(
              batch_count, len(batches), len(batch)
          )
      )
      upload_batch(args, batch, upsert)
    print('   Batch processing complete.')
    batch_count += 1


if __name__ == '__main__':
  main()