atest/logstorage/atest_gcp_utils.py - platform/tools/asuite - Git at Google

 # Copyright (C) 2020 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.
 """Utility functions for atest."""
 from __future__ import print_function

 import getpass
 import logging
 import os
 import pathlib
 from pathlib import Path
 from socket import socket
 import subprocess
 import time
 from typing import Any, Callable
 import uuid

 from atest import atest_utils
 from atest import constants
 from atest.atest_enum import DetectType
 from atest.metrics import metrics
 import httplib2
 from oauth2client import client as oauth2_client
 from oauth2client import contrib as oauth2_contrib
 from oauth2client import tools as oauth2_tools


 class RunFlowFlags:
   """Flags for oauth2client.tools.run_flow."""

   def __init__(self, browser_auth):
     self.auth_host_port = [8080, 8090]
     self.auth_host_name = 'localhost'
     self.logging_level = 'ERROR'
     self.noauth_local_webserver = not browser_auth


 class GCPHelper:
   """GCP bucket helper class."""

   def __init__(
       self,
       client_id=None,
       client_secret=None,
       user_agent=None,
       scope=constants.SCOPE_BUILD_API_SCOPE,
   ):
     """Init stuff for GCPHelper class.

     Args:
         client_id: String, client id from the cloud project.
         client_secret: String, client secret for the client_id.
         user_agent: The user agent for the credential.
         scope: String, scopes separated by space.
     """
     self.client_id = client_id
     self.client_secret = client_secret
     self.user_agent = user_agent
     self.scope = scope

   def get_refreshed_credential_from_file(self, creds_file_path):
     """Get refreshed credential from file.

     Args:
         creds_file_path: Credential file path.

     Returns:
         An oauth2client.OAuth2Credentials instance.
     """
     credential = self.get_credential_from_file(creds_file_path)
     if credential:
       try:
         credential.refresh(httplib2.Http())
       except oauth2_client.AccessTokenRefreshError as e:
         logging.debug('Token refresh error: %s', e)
       if not credential.invalid:
         return credential
     logging.debug('Cannot get credential.')
     return None

   def get_credential_from_file(self, creds_file_path):
     """Get credential from file.

     Args:
         creds_file_path: Credential file path.

     Returns:
         An oauth2client.OAuth2Credentials instance.
     """
     storage = oauth2_contrib.multiprocess_file_storage.get_credential_storage(
         filename=os.path.abspath(creds_file_path),
         client_id=self.client_id,
         user_agent=self.user_agent,
         scope=self.scope,
     )
     return storage.get()

   def get_credential_with_auth_flow(self, creds_file_path):
     """Get Credential object from file.

     Get credential object from file. Run oauth flow if haven't authorized
     before.

     Args:
         creds_file_path: Credential file path.

     Returns:
         An oauth2client.OAuth2Credentials instance.
     """
     credentials = None
     # SSO auth
     try:
       token = self._get_sso_access_token()
       credentials = oauth2_client.AccessTokenCredentials(token, 'atest')
       if credentials:
         return credentials
     # pylint: disable=broad-except
     except Exception as e:
       logging.debug('Exception:%s', e)
     # GCP auth flow
     credentials = self.get_refreshed_credential_from_file(creds_file_path)
     if not credentials:
       storage = oauth2_contrib.multiprocess_file_storage.get_credential_storage(
           filename=os.path.abspath(creds_file_path),
           client_id=self.client_id,
           user_agent=self.user_agent,
           scope=self.scope,
       )
       return self._run_auth_flow(storage)
     return credentials

   def _run_auth_flow(self, storage):
     """Get user oauth2 credentials.

     Using the loopback IP address flow for desktop clients.

     Args:
         storage: GCP storage object.

     Returns:
         An oauth2client.OAuth2Credentials instance.
     """
     flags = RunFlowFlags(browser_auth=True)

     # Get a free port on demand.
     port = None
     while not port or port < 10000:
       with socket() as local_socket:
         local_socket.bind(('', 0))
         _, port = local_socket.getsockname()
     _localhost_port = port
     _direct_uri = f'http://localhost:{_localhost_port}'
     flow = oauth2_client.OAuth2WebServerFlow(
         client_id=self.client_id,
         client_secret=self.client_secret,
         scope=self.scope,
         user_agent=self.user_agent,
         redirect_uri=f'{_direct_uri}',
     )
     credentials = oauth2_tools.run_flow(flow=flow, storage=storage, flags=flags)
     return credentials

   @staticmethod
   def _get_sso_access_token():
     """Use stubby command line to exchange corp sso to a scoped oauth

     token.

     Returns:
         A token string.
     """
     if not constants.TOKEN_EXCHANGE_COMMAND:
       return None

     request = constants.TOKEN_EXCHANGE_REQUEST.format(
         user=getpass.getuser(), scope=constants.SCOPE
     )
     # The output format is: oauth2_token: "<TOKEN>"
     return subprocess.run(
         constants.TOKEN_EXCHANGE_COMMAND,
         input=request,
         check=True,
         text=True,
         shell=True,
         stdout=subprocess.PIPE,
     ).stdout.split('"')[1]


 # TODO: The usage of build_client should be removed from this method because
 # it's not related to this module. For now, we temporarily declare the return
 # type hint for build_client_creator to be Any to avoid circular importing.
 def do_upload_flow(
     extra_args: dict[str, str],
     build_client_creator: Callable,
     invocation_properties: dict[str, str] = None,
 ) -> tuple:
   """Run upload flow.

   Asking user's decision and do the related steps.

   Args:
       extra_args: Dict of extra args to add to test run.
       build_client_creator: A function that takes a credential and returns a
         BuildClient object.
       invocation_properties: Additional invocation properties to write into the
         invocation.

   Return:
       A tuple of credential object and invocation information dict.
   """
   invocation_properties = invocation_properties or {}
   fetch_cred_start = time.time()
   creds = fetch_credential()
   metrics.LocalDetectEvent(
       detect_type=DetectType.FETCH_CRED_MS,
       result=int((time.time() - fetch_cred_start) * 1000),
   )
   if creds:
     prepare_upload_start = time.time()
     build_client = build_client_creator(creds)
     inv, workunit, local_build_id, build_target = _prepare_data(
         build_client, invocation_properties
     )
     metrics.LocalDetectEvent(
         detect_type=DetectType.UPLOAD_PREPARE_MS,
         result=int((time.time() - prepare_upload_start) * 1000),
     )
     extra_args[constants.INVOCATION_ID] = inv['invocationId']
     extra_args[constants.WORKUNIT_ID] = workunit['id']
     extra_args[constants.LOCAL_BUILD_ID] = local_build_id
     extra_args[constants.BUILD_TARGET] = build_target
     if not os.path.exists(os.path.dirname(constants.TOKEN_FILE_PATH)):
       os.makedirs(os.path.dirname(constants.TOKEN_FILE_PATH))
     with open(constants.TOKEN_FILE_PATH, 'w') as token_file:
       if creds.token_response:
         token_file.write(creds.token_response['access_token'])
       else:
         token_file.write(creds.access_token)
     return creds, inv
   return None, None


 def fetch_credential():
   """Fetch the credential object."""
   creds_path = atest_utils.get_config_folder().joinpath(
       constants.CREDENTIAL_FILE_NAME
   )
   return GCPHelper(
       client_id=constants.CLIENT_ID,
       client_secret=constants.CLIENT_SECRET,
       user_agent='atest',
   ).get_credential_with_auth_flow(creds_path)


 def _prepare_data(client, invocation_properties: dict[str, str]):
   """Prepare data for build api using.

   Args:
       build_client: The logstorage_utils.BuildClient object.
       invocation_properties: Additional invocation properties to write into the
         invocation.

   Return:
       invocation and workunit object.
       build id and build target of local build.
   """
   try:
     logging.disable(logging.INFO)
     external_id = str(uuid.uuid4())
     branch = _get_branch(client)
     target = _get_target(branch, client)
     build_record = client.insert_local_build(external_id, target, branch)
     client.insert_build_attempts(build_record)
     invocation = client.insert_invocation(build_record, invocation_properties)
     workunit = client.insert_work_unit(invocation)
     return invocation, workunit, build_record['buildId'], target
   finally:
     logging.disable(logging.NOTSET)


 def _get_branch(build_client):
   """Get source code tree branch.

   Args:
       build_client: The build client object.

   Return:
       "git_main" in internal git, "aosp-main" otherwise.
   """
   default_branch = 'git_main' if constants.CREDENTIAL_FILE_NAME else 'aosp-main'
   local_branch = 'git_%s' % atest_utils.get_manifest_branch()
   branch = build_client.get_branch(local_branch)
   return local_branch if branch else default_branch


 def _get_target(branch, build_client):
   """Get local build selected target.

   Args:
       branch: The branch want to check.
       build_client: The build client object.

   Return:
       The matched build target, "aosp_x86_64-trunk_staging-userdebug"
       otherwise.
   """
   default_target = 'aosp_x86_64-trunk_staging-userdebug'
   local_target = atest_utils.get_build_target()
   targets = [t['target'] for t in build_client.list_target(branch)['targets']]
   return local_target if local_target in targets else default_target
	# Copyright (C) 2020 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.
	"""Utility functions for atest."""
	from __future__ import print_function

	import getpass
	import logging
	import os
	import pathlib
	from pathlib import Path
	from socket import socket
	import subprocess
	import time
	from typing import Any, Callable
	import uuid

	from atest import atest_utils
	from atest import constants
	from atest.atest_enum import DetectType
	from atest.metrics import metrics
	import httplib2
	from oauth2client import client as oauth2_client
	from oauth2client import contrib as oauth2_contrib
	from oauth2client import tools as oauth2_tools


	class RunFlowFlags:
	"""Flags for oauth2client.tools.run_flow."""

	def __init__(self, browser_auth):
	self.auth_host_port = [8080, 8090]
	self.auth_host_name = 'localhost'
	self.logging_level = 'ERROR'
	self.noauth_local_webserver = not browser_auth


	class GCPHelper:
	"""GCP bucket helper class."""

	def __init__(
	self,
	client_id=None,
	client_secret=None,
	user_agent=None,
	scope=constants.SCOPE_BUILD_API_SCOPE,
	):
	"""Init stuff for GCPHelper class.

	Args:
	client_id: String, client id from the cloud project.
	client_secret: String, client secret for the client_id.
	user_agent: The user agent for the credential.
	scope: String, scopes separated by space.
	"""
	self.client_id = client_id
	self.client_secret = client_secret
	self.user_agent = user_agent
	self.scope = scope

	def get_refreshed_credential_from_file(self, creds_file_path):
	"""Get refreshed credential from file.

	Args:
	creds_file_path: Credential file path.

	Returns:
	An oauth2client.OAuth2Credentials instance.
	"""
	credential = self.get_credential_from_file(creds_file_path)
	if credential:
	try:
	credential.refresh(httplib2.Http())
	except oauth2_client.AccessTokenRefreshError as e:
	logging.debug('Token refresh error: %s', e)
	if not credential.invalid:
	return credential
	logging.debug('Cannot get credential.')
	return None

	def get_credential_from_file(self, creds_file_path):
	"""Get credential from file.

	Args:
	creds_file_path: Credential file path.

	Returns:
	An oauth2client.OAuth2Credentials instance.
	"""
	storage = oauth2_contrib.multiprocess_file_storage.get_credential_storage(
	filename=os.path.abspath(creds_file_path),
	client_id=self.client_id,
	user_agent=self.user_agent,
	scope=self.scope,
	)
	return storage.get()

	def get_credential_with_auth_flow(self, creds_file_path):
	"""Get Credential object from file.

	Get credential object from file. Run oauth flow if haven't authorized
	before.

	Args:
	creds_file_path: Credential file path.

	Returns:
	An oauth2client.OAuth2Credentials instance.
	"""
	credentials = None
	# SSO auth
	try:
	token = self._get_sso_access_token()
	credentials = oauth2_client.AccessTokenCredentials(token, 'atest')
	if credentials:
	return credentials
	# pylint: disable=broad-except
	except Exception as e:
	logging.debug('Exception:%s', e)
	# GCP auth flow
	credentials = self.get_refreshed_credential_from_file(creds_file_path)
	if not credentials:
	storage = oauth2_contrib.multiprocess_file_storage.get_credential_storage(
	filename=os.path.abspath(creds_file_path),
	client_id=self.client_id,
	user_agent=self.user_agent,
	scope=self.scope,
	)
	return self._run_auth_flow(storage)
	return credentials

	def _run_auth_flow(self, storage):
	"""Get user oauth2 credentials.

	Using the loopback IP address flow for desktop clients.

	Args:
	storage: GCP storage object.

	Returns:
	An oauth2client.OAuth2Credentials instance.
	"""
	flags = RunFlowFlags(browser_auth=True)

	# Get a free port on demand.
	port = None
	while not port or port < 10000:
	with socket() as local_socket:
	local_socket.bind(('', 0))
	_, port = local_socket.getsockname()
	_localhost_port = port
	_direct_uri = f'http://localhost:{_localhost_port}'
	flow = oauth2_client.OAuth2WebServerFlow(
	client_id=self.client_id,
	client_secret=self.client_secret,
	scope=self.scope,
	user_agent=self.user_agent,
	redirect_uri=f'{_direct_uri}',
	)
	credentials = oauth2_tools.run_flow(flow=flow, storage=storage, flags=flags)
	return credentials

	@staticmethod
	def _get_sso_access_token():
	"""Use stubby command line to exchange corp sso to a scoped oauth

	token.

	Returns:
	A token string.
	"""
	if not constants.TOKEN_EXCHANGE_COMMAND:
	return None

	request = constants.TOKEN_EXCHANGE_REQUEST.format(
	user=getpass.getuser(), scope=constants.SCOPE
	)
	# The output format is: oauth2_token: "<TOKEN>"
	return subprocess.run(
	constants.TOKEN_EXCHANGE_COMMAND,
	input=request,
	check=True,
	text=True,
	shell=True,
	stdout=subprocess.PIPE,
	).stdout.split('"')[1]


	# TODO: The usage of build_client should be removed from this method because
	# it's not related to this module. For now, we temporarily declare the return
	# type hint for build_client_creator to be Any to avoid circular importing.
	def do_upload_flow(
	extra_args: dict[str, str],
	build_client_creator: Callable,
	invocation_properties: dict[str, str] = None,
	) -> tuple:
	"""Run upload flow.

	Asking user's decision and do the related steps.

	Args:
	extra_args: Dict of extra args to add to test run.
	build_client_creator: A function that takes a credential and returns a
	BuildClient object.
	invocation_properties: Additional invocation properties to write into the
	invocation.

	Return:
	A tuple of credential object and invocation information dict.
	"""
	invocation_properties = invocation_properties or {}
	fetch_cred_start = time.time()
	creds = fetch_credential()
	metrics.LocalDetectEvent(
	detect_type=DetectType.FETCH_CRED_MS,
	result=int((time.time() - fetch_cred_start) * 1000),
	)
	if creds:
	prepare_upload_start = time.time()
	build_client = build_client_creator(creds)
	inv, workunit, local_build_id, build_target = _prepare_data(
	build_client, invocation_properties
	)
	metrics.LocalDetectEvent(
	detect_type=DetectType.UPLOAD_PREPARE_MS,
	result=int((time.time() - prepare_upload_start) * 1000),
	)
	extra_args[constants.INVOCATION_ID] = inv['invocationId']
	extra_args[constants.WORKUNIT_ID] = workunit['id']
	extra_args[constants.LOCAL_BUILD_ID] = local_build_id
	extra_args[constants.BUILD_TARGET] = build_target
	if not os.path.exists(os.path.dirname(constants.TOKEN_FILE_PATH)):
	os.makedirs(os.path.dirname(constants.TOKEN_FILE_PATH))
	with open(constants.TOKEN_FILE_PATH, 'w') as token_file:
	if creds.token_response:
	token_file.write(creds.token_response['access_token'])
	else:
	token_file.write(creds.access_token)
	return creds, inv
	return None, None


	def fetch_credential():
	"""Fetch the credential object."""
	creds_path = atest_utils.get_config_folder().joinpath(
	constants.CREDENTIAL_FILE_NAME
	)
	return GCPHelper(
	client_id=constants.CLIENT_ID,
	client_secret=constants.CLIENT_SECRET,
	user_agent='atest',
	).get_credential_with_auth_flow(creds_path)


	def _prepare_data(client, invocation_properties: dict[str, str]):
	"""Prepare data for build api using.

	Args:
	build_client: The logstorage_utils.BuildClient object.
	invocation_properties: Additional invocation properties to write into the
	invocation.

	Return:
	invocation and workunit object.
	build id and build target of local build.
	"""
	try:
	logging.disable(logging.INFO)
	external_id = str(uuid.uuid4())
	branch = _get_branch(client)
	target = _get_target(branch, client)
	build_record = client.insert_local_build(external_id, target, branch)
	client.insert_build_attempts(build_record)
	invocation = client.insert_invocation(build_record, invocation_properties)
	workunit = client.insert_work_unit(invocation)
	return invocation, workunit, build_record['buildId'], target
	finally:
	logging.disable(logging.NOTSET)


	def _get_branch(build_client):
	"""Get source code tree branch.

	Args:
	build_client: The build client object.

	Return:
	"git_main" in internal git, "aosp-main" otherwise.
	"""
	default_branch = 'git_main' if constants.CREDENTIAL_FILE_NAME else 'aosp-main'
	local_branch = 'git_%s' % atest_utils.get_manifest_branch()
	branch = build_client.get_branch(local_branch)
	return local_branch if branch else default_branch


	def _get_target(branch, build_client):
	"""Get local build selected target.

	Args:
	branch: The branch want to check.
	build_client: The build client object.

	Return:
	The matched build target, "aosp_x86_64-trunk_staging-userdebug"
	otherwise.
	"""
	default_target = 'aosp_x86_64-trunk_staging-userdebug'
	local_target = atest_utils.get_build_target()
	targets = [t['target'] for t in build_client.list_target(branch)['targets']]
	return local_target if local_target in targets else default_target