Use oauth token from LOAS/SSO when querying build api am: f76785651c

Original change: https://android-review.googlesource.com/c/toolchain/pgo-profiles/+/1683347

Change-Id: I0e88c33f118231cd55a28879ec24fc4bdc0f81a4
diff --git a/scripts/android_build_client.py b/scripts/android_build_client.py
index da39262..7cee477 100644
--- a/scripts/android_build_client.py
+++ b/scripts/android_build_client.py
@@ -1,4 +1,4 @@
-#!/usr/bin/env python
+#!/usr/bin/env python3
 #
 # Copyright (C) 2019 The Android Open Source Project
 #
@@ -20,32 +20,59 @@
     import apiclient.http
     from oauth2client import client as oauth2_client
 except ImportError:
-    missingImportString = '''
+    missingImportString = """
   Missing necessary libraries. Try doing the following:
-  $ sudo apt-get install python-pip
-  $ sudo pip install --upgrade google-api-python-client
-  $ sudo pip install --upgrade oauth2client
-'''
+  $ sudo apt-get install python3-pip
+  $ pip install --user --upgrade google-api-python-client
+  $ pip install --user --upgrade oauth2client
+"""
     raise ImportError(missingImportString)
 
 import io
+import getpass
 import os
 
+import utils
+
 ANDROID_BUILD_API_SCOPE = (
     'https://www.googleapis.com/auth/androidbuild.internal')
 ANDROID_BUILD_API_NAME = 'androidbuildinternal'
 ANDROID_BUILD_API_VERSION = 'v2beta1'
-TRADEFED_KEY_FILE = '/google/data/ro/teams/tradefed/configs/tradefed.json'
 CHUNK_SIZE = 10 * 1024 * 1024  # 10M
 
 ANDROID_PGO_BUILD = 'pgo-coral-config1'
 
+STUBBY_COMMAND_PATH = '/google/data/ro/teams/android-llvm/tests/sso_stubby_cmd.sh'
+STUBBY_REQUEST = """
+target: {{
+  scope: GAIA_USER
+  name: "{user}@google.com"
+}}
+target_credential: {{
+  type: OAUTH2_TOKEN
+  oauth2_attributes: {{
+    scope: '{scope}'
+  }}
+}}
+"""
+
+
+def _get_oauth2_token():
+    request = STUBBY_REQUEST.format(
+        user=getpass.getuser(), scope=ANDROID_BUILD_API_SCOPE)
+    with open(STUBBY_COMMAND_PATH) as stubby_command_file:
+        stubby_command = stubby_command_file.read().strip().split()
+    output = utils.check_output(stubby_command, input=request)
+    # output is of the format:
+    # oauth2_token: "<TOKEN>"
+    return output.split('"')[1]
+
 
 class AndroidBuildClient(object):
 
     def __init__(self):
-        creds = oauth2_client.GoogleCredentials.from_stream(TRADEFED_KEY_FILE)
-        self.creds = creds.create_scoped([ANDROID_BUILD_API_SCOPE])
+        creds = oauth2_client.AccessTokenCredentials(
+            access_token=_get_oauth2_token(), user_agent='unused/1.0')
 
         self.client = apiclient.discovery.build(
             ANDROID_BUILD_API_NAME,
diff --git a/scripts/update_profiles.py b/scripts/update_profiles.py
index 1413af1..375e96d 100755
--- a/scripts/update_profiles.py
+++ b/scripts/update_profiles.py
@@ -1,4 +1,4 @@
-#!/usr/bin/env python
+#!/usr/bin/env python3
 #
 # Copyright (C) 2019 The Android Open Source Project
 #
@@ -172,6 +172,7 @@
 
 def main():
     args = parse_args()
+    utils.check_gcertstatus()
 
     if args.benchmark == 'ALL':
         worklist = KNOWN_BENCHMARKS[1:]
diff --git a/scripts/utils.py b/scripts/utils.py
index 5f39c06..035c1b8 100644
--- a/scripts/utils.py
+++ b/scripts/utils.py
@@ -42,19 +42,30 @@
     subprocess.check_call(cmd, *args, **kwargs)
 
 
+def check_output(cmd, *args, **kwargs):
+    """subprocess.check_output with logging."""
+    logger().info('check_output: %s', subprocess.list2cmdline(cmd))
+    return subprocess.run(
+        cmd, *args, **kwargs, check=True, text=True,
+        stdout=subprocess.PIPE).stdout
+
+
 def android_build_top():
     return os.path.realpath(os.path.join(THIS_DIR, '../../..'))
 
 
 def clang_build():
-    gofile = os.path.join(android_build_top(), 'build', 'soong', 'cc', 'config', 'global.go')
+    gofilename = os.path.join(android_build_top(), 'build', 'soong', 'cc',
+                              'config', 'global.go')
     try:
-        lines = file(gofile).readlines()
+        with open(gofilename) as gofile:
+            lines = gofile.readlines()
         versionLine = [l for l in lines if 'ClangDefaultVersion' in l][0]
         start, end = versionLine.index('"'), versionLine.rindex('"')
-        return versionLine[start + 1: end]
+        return versionLine[start + 1:end]
     except Exception as err:
-        raise RuntimeError("Extracting Clang version failed with {0}".format(err))
+        raise RuntimeError(
+            'Extracting Clang version failed with {0}'.format(err))
 
 
 def llvm_profdata():
@@ -64,3 +75,12 @@
 
 def run_llvm_profdata(inputs, output):
     check_call([llvm_profdata(), 'merge', '-output=' + output] + inputs)
+
+
+def check_gcertstatus():
+    """Ensure gcert valid for > 1 hour."""
+    try:
+        check_call(['gcertstatus', '-quiet', '-check_remaining=1h'])
+    except subprocess.CalledProcessError:
+        print('Run prodaccess before executing this script.')
+        raise