xref: /aosp_15_r20/external/autotest/utils/frozen_chromite/lib/cipd.py (revision 9c5db1993ded3edbeafc8092d69fe5de2ee02df7)
1# -*- coding: utf-8 -*-
2# Copyright 2015 The Chromium OS Authors. All rights reserved.
3# Use of this source code is governed by a BSD-style license that can be
4# found in the LICENSE file.
5
6"""Module to download and run the CIPD client.
7
8CIPD is the Chrome Infra Package Deployer, a simple method of resolving a
9package/version into a GStorage link and installing them.
10"""
11
12from __future__ import print_function
13
14import hashlib
15import json
16import os
17import pprint
18import tempfile
19
20import httplib2
21from six.moves import urllib
22
23import autotest_lib.utils.frozen_chromite.lib.cros_logging as log
24from autotest_lib.utils.frozen_chromite.lib import cache
25from autotest_lib.utils.frozen_chromite.lib import osutils
26from autotest_lib.utils.frozen_chromite.lib import path_util
27from autotest_lib.utils.frozen_chromite.lib import cros_build_lib
28from autotest_lib.utils.frozen_chromite.utils import memoize
29
30# pylint: disable=line-too-long
31# CIPD client to download.
32#
33# This is version "git_revision:db7a486094873e3944b8e27ab5b23a3ae3c401e7".
34#
35# To switch to another version:
36#   1. Find it in CIPD Web UI, e.g.
37#      https://chrome-infra-packages.appspot.com/p/infra/tools/cipd/linux-amd64/+/latest
38#   2. Look up SHA256 there.
39# pylint: enable=line-too-long
40CIPD_CLIENT_PACKAGE = 'infra/tools/cipd/linux-amd64'
41CIPD_CLIENT_SHA256 = (
42    'ea6b7547ddd316f32fd9974f598949c3f8f22f6beb8c260370242d0d84825162')
43
44CHROME_INFRA_PACKAGES_API_BASE = (
45    'https://chrome-infra-packages.appspot.com/prpc/cipd.Repository/')
46
47
48class Error(Exception):
49  """Raised on fatal errors."""
50
51
52def _ChromeInfraRequest(method, request):
53  """Makes a request to the Chrome Infra Packages API with httplib2.
54
55  Args:
56    method: Name of RPC method to call.
57    request: RPC request body.
58
59  Returns:
60    Deserialized RPC response body.
61  """
62  resp, body = httplib2.Http().request(
63      uri=CHROME_INFRA_PACKAGES_API_BASE+method,
64      method='POST',
65      headers={
66          'Accept': 'application/json',
67          'Content-Type': 'application/json',
68          'User-Agent': 'chromite',
69      },
70      body=json.dumps(request))
71  if resp.status != 200:
72    raise Error('Got HTTP %d from CIPD %r: %s' % (resp.status, method, body))
73  try:
74    return json.loads(body.lstrip(b")]}'\n"))
75  except ValueError:
76    raise Error('Bad response from CIPD server:\n%s' % (body,))
77
78
79def _DownloadCIPD(instance_sha256):
80  """Finds the CIPD download link and requests the binary.
81
82  Args:
83    instance_sha256: The version of CIPD client to download.
84
85  Returns:
86    The CIPD binary as a string.
87  """
88  # Grab the signed URL to fetch the client binary from.
89  resp = _ChromeInfraRequest('DescribeClient', {
90      'package': CIPD_CLIENT_PACKAGE,
91      'instance': {
92          'hashAlgo': 'SHA256',
93          'hexDigest': instance_sha256,
94      },
95  })
96  if 'clientBinary' not in resp:
97    log.error(
98        'Error requesting the link to download CIPD from. Got:\n%s',
99        pprint.pformat(resp))
100    raise Error('Failed to bootstrap CIPD client')
101
102  # Download the actual binary.
103  http = httplib2.Http(cache=None)
104  response, binary = http.request(uri=resp['clientBinary']['signedUrl'])
105  if response.status != 200:
106    raise Error('Got a %d response from Google Storage.' % response.status)
107
108  # Check SHA256 matches what server expects.
109  digest = hashlib.sha256(binary).hexdigest()
110  for alias in resp['clientRefAliases']:
111    if alias['hashAlgo'] == 'SHA256':
112      if digest != alias['hexDigest']:
113        raise Error(
114            'Unexpected CIPD client SHA256: got %s, want %s' %
115            (digest, alias['hexDigest']))
116      break
117  else:
118    raise Error("CIPD server didn't provide expected SHA256")
119
120  return binary
121
122
123class CipdCache(cache.RemoteCache):
124  """Supports caching of the CIPD download."""
125  def _Fetch(self, url, local_path):
126    instance_sha256 = urllib.parse.urlparse(url).netloc
127    binary = _DownloadCIPD(instance_sha256)
128    log.info('Fetched CIPD package %s:%s', CIPD_CLIENT_PACKAGE, instance_sha256)
129    osutils.WriteFile(local_path, binary, mode='wb')
130    os.chmod(local_path, 0o755)
131
132
133def GetCIPDFromCache():
134  """Checks the cache, downloading CIPD if it is missing.
135
136  Returns:
137    Path to the CIPD binary.
138  """
139  cache_dir = os.path.join(path_util.GetCacheDir(), 'cipd')
140  bin_cache = CipdCache(cache_dir)
141  key = (CIPD_CLIENT_SHA256,)
142  ref = bin_cache.Lookup(key)
143  ref.SetDefault('cipd://' + CIPD_CLIENT_SHA256)
144  return ref.path
145
146
147def GetInstanceID(cipd_path, package, version, service_account_json=None):
148  """Get the latest instance ID for ref latest.
149
150  Args:
151    cipd_path: The path to a cipd executable. GetCIPDFromCache can give this.
152    package: A string package name.
153    version: A string version of package.
154    service_account_json: The path of the service account credentials.
155
156  Returns:
157    A string instance ID.
158  """
159  service_account_flag = []
160  if service_account_json:
161    service_account_flag = ['-service-account-json', service_account_json]
162
163  result = cros_build_lib.run(
164      [cipd_path, 'resolve', package, '-version', version] +
165      service_account_flag, capture_output=True, encoding='utf-8')
166  # An example output of resolve is like:
167  #   Packages:\n package:instance_id
168  return result.output.splitlines()[-1].split(':')[-1]
169
170
171@memoize.Memoize
172def InstallPackage(cipd_path, package, instance_id, destination,
173                   service_account_json=None):
174  """Installs a package at a given destination using cipd.
175
176  Args:
177    cipd_path: The path to a cipd executable. GetCIPDFromCache can give this.
178    package: A package name.
179    instance_id: The version of the package to install.
180    destination: The folder to install the package under.
181    service_account_json: The path of the service account credentials.
182
183  Returns:
184    The path of the package.
185  """
186  destination = os.path.join(destination, package)
187
188  service_account_flag = []
189  if service_account_json:
190    service_account_flag = ['-service-account-json', service_account_json]
191
192  with tempfile.NamedTemporaryFile() as f:
193    f.write(('%s %s' % (package, instance_id)).encode('utf-8'))
194    f.flush()
195
196    cros_build_lib.run(
197        [cipd_path, 'ensure', '-root', destination, '-list', f.name]
198        + service_account_flag,
199        capture_output=True)
200
201  return destination
202
203
204def CreatePackage(cipd_path, package, in_dir, tags, refs,
205                  cred_path=None):
206  """Create (build and register) a package using cipd.
207
208  Args:
209    cipd_path: The path to a cipd executable. GetCIPDFromCache can give this.
210    package: A package name.
211    in_dir: The directory to create the package from.
212    tags: A mapping of tags to apply to the package.
213    refs: An Iterable of refs to apply to the package.
214    cred_path: The path of the service account credentials.
215  """
216  args = [
217      cipd_path, 'create',
218      '-name', package,
219      '-in', in_dir,
220  ]
221  for key, value in tags.items():
222    args.extend(['-tag', '%s:%s' % (key, value)])
223  for ref in refs:
224    args.extend(['-ref', ref])
225  if cred_path:
226    args.extend(['-service-account-json', cred_path])
227
228  cros_build_lib.run(args, capture_output=True)
229
230
231def BuildPackage(cipd_path, package, in_dir, outfile):
232  """Build a package using cipd.
233
234  Args:
235    cipd_path: The path to a cipd executable. GetCIPDFromCache can give this.
236    package: A package name.
237    in_dir: The directory to create the package from.
238    outfile: Output file.  Should have extension .cipd
239  """
240  args = [
241      cipd_path, 'pkg-build',
242      '-name', package,
243      '-in', in_dir,
244      '-out', outfile,
245  ]
246  cros_build_lib.run(args, capture_output=True)
247
248
249def RegisterPackage(cipd_path, package_file, tags, refs, cred_path=None):
250  """Register and upload a package using cipd.
251
252  Args:
253    cipd_path: The path to a cipd executable. GetCIPDFromCache can give this.
254    package_file: The path to a .cipd package file.
255    tags: A mapping of tags to apply to the package.
256    refs: An Iterable of refs to apply to the package.
257    cred_path: The path of the service account credentials.
258  """
259  args = [cipd_path, 'pkg-register', package_file]
260  for key, value in tags.items():
261    args.extend(['-tag', '%s:%s' % (key, value)])
262  for ref in refs:
263    args.extend(['-ref', ref])
264  if cred_path:
265    args.extend(['-service-account-json', cred_path])
266  cros_build_lib.run(args, capture_output=True)
267