lib/auth.py

# Copyright 2015 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Functions for authenticating httplib2 requests with OAuth2 tokens."""

from __future__ import print_function

import functools
import os
import sys

import chromite.lib.cros_logging as log
from chromite.lib import cipd
from chromite.lib import cros_build_lib
from chromite.lib import retry_util
from chromite.lib import path_util


REFRESH_STATUS_CODES = [401]


def _GetAuthUtil(instance_id='latest'):
  """Returns a path to the authutil binary.

  This will download and install the authutil package if it is not already
  deployed.

  Args:
    instance_id: The instance-id of the package to install. Defaults to 'latest'

  Returns:
    the path to the authutil binary.
  """
  cache_dir = os.path.join(path_util.GetCacheDir(), 'cipd/packages')
  path = cipd.InstallPackage(
      cipd.GetCIPDFromCache(),
      'infra/tools/authutil/linux-amd64',
      instance_id,
      destination=cache_dir)

  return os.path.join(path, 'authutil')


def _Login():
  """Logs a user into chrome-infra-auth using authutil.

  Runs 'authutil login' to get a OAuth2 refresh token.

  Returns:
    Whether the login process was successful.
  """
  log.info('Logging into chrome-infra-auth...')
  result = retry_util.RunCommandWithRetries(
      3,
      args=[[_GetAuthUtil(), 'login']],
      kwargs={'error_code_ok': True})
  if not result.returncode:
    return True
  print(result.error, file=sys.stderr)
  log.error("Error logging in to chrome-infra-auth.")
  return False


def GetAccessToken(service_account_json=None):
  """Returns an OAuth2 access token using authutil.

  Attempts to get a token using 'authutil token'. If that fails, retries after
  running 'authutil login'.

  Args:
    service_account_json: A optional path to a service account to use

  Returns:
    The access token string.
  """
  flags = []
  if service_account_json:
    flags += ['-service_account_json', service_account_json]
  result = cros_build_lib.RunCommand(
      [_GetAuthUtil(), 'token'] + flags,
      capture_output=True,
      error_code_ok=True)

  if result.returncode:
    if not _Login():
      return
    return GetAccessToken()
  else:
    return result.output.strip()


def Authorize(get_access_token, http):
  """Monkey patches authentication logic of httplib2.Http instance.

  The modified http.request method will add authentication headers to each
  request and will refresh access_tokens when a 401 is received on a
  request.

  Args:
    get_access_token: A function which returns an access token.
    http: An instance of httplib2.Http.

  Returns:
    A modified instance of http that was passed in.
  """
  # Adapted from oauth2client.OAuth2Credentials.authorize.
  # We can't use oauthclient2 because the import will fail on slaves due to
  # missing PyOpenSSL (crbug.com/498467).
  request_orig = http.request

  @functools.wraps(request_orig)
  def new_request(*args, **kwargs):
    headers = kwargs.get('headers', {}).copy()
    headers['Authorization'] = 'Bearer %s' % get_access_token()
    kwargs['headers'] = headers

    resp, content = request_orig(*args, **kwargs)
    if resp.status in REFRESH_STATUS_CODES:
      log.info('Refreshing due to a %s', resp.status)
      # TODO(phobbs): delete the "access_token" key from the token file used.
      headers['Authorization'] = 'Bearer %s' % get_access_token().token
      resp, content = request_orig(*args, **kwargs)

    return resp, content

  http.request = new_request
  return http