xref: /aosp_15_r20/external/perfetto/infra/ci/common_utils.py (revision 6dbdd20afdafa5e3ca9b8809fa73465d530080dc)
1# Copyright (C) 2019 The Android Open Source Project
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#      http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15import asyncio
16import concurrent.futures
17import google.auth
18import google.auth.transport.requests
19import json
20import logging
21import os
22import requests
23
24from base64 import b64encode
25from datetime import datetime
26from config import PROJECT
27
28# Thread pool for making http requests asynchronosly.
29thread_pool = concurrent.futures.ThreadPoolExecutor(max_workers=8)
30
31# Caller has to initialize this
32SCOPES = []
33cached_gerrit_creds = None
34cached_oauth2_creds = None
35
36
37class ConcurrentModificationError(Exception):
38  pass
39
40
41def get_access_token():
42  global cached_oauth2_creds
43  creds = cached_oauth2_creds
44  if creds is None or not creds.valid or creds.expired:
45    creds, _project = google.auth.default(scopes=SCOPES)
46    request = google.auth.transport.requests.Request()
47    creds.refresh(request)
48    cached_oauth2_creds = creds
49  return creds.token
50
51
52def get_gerrit_credentials():
53  '''Retrieve the credentials used to authenticate Gerrit requests
54
55  Returns a tuple (user, gitcookie). These fields are obtained from the Gerrit
56  'New HTTP password' page which generates a .gitcookie file and stored in the
57  project datastore.
58  user: typically looks like git-user.gmail.com.
59  gitcookie: is the password after the = token.
60  '''
61  global cached_gerrit_creds
62  if cached_gerrit_creds is None:
63    body = {'query': {'kind': [{'name': 'GerritAuth'}]}}
64    res = req(
65        'POST',
66        'https://datastore.googleapis.com/v1/projects/%s:runQuery' % PROJECT,
67        body=body)
68    auth = res['batch']['entityResults'][0]['entity']['properties']
69    user = auth['user']['stringValue']
70    gitcookie = auth['gitcookie']['stringValue']
71    cached_gerrit_creds = user, gitcookie
72  return cached_gerrit_creds
73
74
75async def req_async(method, url, body=None, gerrit=False):
76  loop = asyncio.get_running_loop()
77  # run_in_executor cannot take kwargs, we need to stick with order.
78  return await loop.run_in_executor(thread_pool, req, method, url, body, gerrit,
79                                    False, None)
80
81
82def req(method, url, body=None, gerrit=False, req_etag=False, etag=None):
83  '''Helper function to handle authenticated HTTP requests.
84
85  Cloud API and Gerrit require two different types of authentication and as
86  such need to be handled differently. The HTTP connection is cached in the
87  TLS slot to avoid refreshing oauth tokens too often for back-to-back requests.
88  Appengine takes care of clearing the TLS slot upon each frontend request so
89  these connections won't be recycled for too long.
90  '''
91  hdr = {'Content-Type': 'application/json; charset=UTF-8'}
92  if gerrit:
93    creds = '%s:%s' % get_gerrit_credentials()
94    auth_header = 'Basic ' + b64encode(creds.encode('utf-8')).decode('utf-8')
95  elif SCOPES:
96    auth_header = 'Bearer ' + get_access_token()
97  logging.debug('%s %s [gerrit=%d]', method, url, gerrit)
98  hdr['Authorization'] = auth_header
99  if req_etag:
100    hdr['X-Firebase-ETag'] = 'true'
101  if etag:
102    hdr['if-match'] = etag
103  body = None if body is None else json.dumps(body)
104  resp = requests.request(method, url, headers=hdr, data=body, timeout=60)
105  res = resp.content.decode('utf-8')
106  resp_etag = resp.headers.get('etag')
107  if resp.status_code == 200:
108    # [4:] is to strip Gerrit XSSI projection prefix.
109    res = json.loads(res[4:] if gerrit else res)
110    return (res, resp_etag) if req_etag else res
111  elif resp.status_code == 412:
112    raise ConcurrentModificationError()
113  else:
114    raise Exception(resp, res)
115
116
117# Datetime functions to deal with the fact that Javascript expects a trailing
118# 'Z' (Z == 'Zulu' == UTC) for timestamps.
119def parse_iso_time(time_str):
120  return datetime.strptime(time_str, r'%Y-%m-%dT%H:%M:%SZ')
121
122
123def utc_now_iso(utcnow=None):
124  return (utcnow or datetime.utcnow()).strftime(r'%Y-%m-%dT%H:%M:%SZ')
125
126
127def defer(coro):
128  loop = asyncio.get_event_loop()
129  task = loop.create_task(coro)
130  task.set_name(coro.cr_code.co_name)
131  return task
132
133
134def init_logging():
135  logging.basicConfig(
136      format='%(levelname)-8s %(asctime)s %(message)s',
137      level=logging.DEBUG if os.getenv('VERBOSE') else logging.INFO,
138      datefmt=r'%Y-%m-%d %H:%M:%S')
139