xref: /aosp_15_r20/external/pigweed/pw_env_setup/py/pw_env_setup/cipd_setup/wrapper.py (revision 61c4878ac05f98d0ceed94b57d316916de578985) 
1#!/usr/bin/env python
2# Copyright 2020 The Pigweed Authors
3#
4# Licensed under the Apache License, Version 2.0 (the "License"); you may not
5# use this file except in compliance with the License. You may obtain a copy of
6# the License at
7#
8#     https://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
12# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
13# License for the specific language governing permissions and limitations under
14# the License.
15"""Installs and then runs cipd.
16
17This script installs cipd in ./tools/ (if necessary) and then executes it,
18passing through all arguments.
19
20Must be tested with Python 2 and Python 3.
21"""
22
23import hashlib
24import os
25import platform
26import ssl
27import subprocess
28import sys
29import base64
30
31try:
32    import httplib  # type: ignore
33except ImportError:
34    import http.client as httplib  # type: ignore[no-redef]
35
36try:
37    import urlparse  # type: ignore
38except ImportError:
39    import urllib.parse as urlparse  # type: ignore[no-redef]
40
41# Generated from the following command. May need to be periodically rerun.
42# $ cipd ls infra/tools/cipd | perl -pe "s[.*/][];s/^/    '/;s/\s*$/',\n/;"
43SUPPORTED_PLATFORMS = (
44    'aix-ppc64',
45    'linux-386',
46    'linux-amd64',
47    'linux-arm64',
48    'linux-armv6l',
49    'linux-mips64',
50    'linux-mips64le',
51    'linux-mipsle',
52    'linux-ppc64',
53    'linux-ppc64le',
54    'linux-s390x',
55    'mac-amd64',
56    'mac-arm64',
57    'windows-386',
58    'windows-amd64',
59)
60
61
62class UnsupportedPlatform(Exception):
63    pass
64
65
66try:
67    SCRIPT_DIR = os.path.dirname(__file__)
68except NameError:  # __file__ not defined.
69    try:
70        SCRIPT_DIR = os.path.join(
71            os.environ['PW_ROOT'],
72            'pw_env_setup',
73            'py',
74            'pw_env_setup',
75            'cipd_setup',
76        )
77    except KeyError:
78        raise Exception('Environment variable PW_ROOT not set')
79
80VERSION_FILE = os.path.join(SCRIPT_DIR, '.cipd_version')
81DIGESTS_FILE = VERSION_FILE + '.digests'
82
83# Put CIPD client in tools so that users can easily get it in their PATH.
84CIPD_HOST = 'chrome-infra-packages.appspot.com'
85
86try:
87    PW_ROOT = os.environ['PW_ROOT']
88except KeyError:
89    try:
90        with open(os.devnull, 'w') as outs:
91            PW_ROOT = (
92                subprocess.check_output(
93                    ['git', 'rev-parse', '--show-toplevel'],
94                    stderr=outs,
95                )
96                .strip()
97                .decode('utf-8')
98            )
99    except subprocess.CalledProcessError:
100        PW_ROOT = ''
101
102# Get default install dir from environment since args cannot always be passed
103# through this script (args are passed as-is to cipd).
104if 'CIPD_PY_INSTALL_DIR' in os.environ:
105    DEFAULT_INSTALL_DIR = os.environ['CIPD_PY_INSTALL_DIR']
106elif PW_ROOT:
107    DEFAULT_INSTALL_DIR = os.path.join(PW_ROOT, '.cipd')
108else:
109    DEFAULT_INSTALL_DIR = ''
110
111
112def platform_normalized():
113    """Normalize platform into format expected in CIPD paths."""
114
115    try:
116        os_name = platform.system().lower()
117        return {
118            'linux': 'linux',
119            'mac': 'mac',
120            'darwin': 'mac',
121            'windows': 'windows',
122        }[os_name]
123    except KeyError:
124        raise Exception('unrecognized os: {}'.format(os_name))
125
126
127def arch_normalized(rosetta=False):
128    """Normalize arch into format expected in CIPD paths."""
129
130    machine = platform.machine()
131    if platform_normalized() == 'mac' and rosetta:
132        return 'amd64'
133    if machine.startswith(('arm', 'aarch')):
134        machine = machine.replace('aarch', 'arm')
135        if machine == 'arm64':
136            return machine
137        return 'armv6l'
138    if machine.endswith('64'):
139        return 'amd64'
140    if machine.endswith('86'):
141        return '386'
142    raise Exception('unrecognized arch: {}'.format(machine))
143
144
145def platform_arch_normalized(rosetta=False):
146    return '{}-{}'.format(platform_normalized(), arch_normalized(rosetta))
147
148
149def user_agent():
150    """Generate a user-agent based on the project name and current hash."""
151
152    try:
153        with open(os.devnull, 'w') as devnull:
154            rev = subprocess.check_output(
155                ['git', '-C', SCRIPT_DIR, 'rev-parse', 'HEAD'],
156                stderr=devnull,
157            ).strip()
158    except subprocess.CalledProcessError:
159        rev = '???'
160
161    if isinstance(rev, bytes):
162        rev = rev.decode()
163
164    return 'pigweed-infra/tools/{}'.format(rev)
165
166
167def actual_hash(path):
168    """Hash the file at path and return it."""
169
170    hasher = hashlib.sha256()
171    with open(path, 'rb') as ins:
172        hasher.update(ins.read())
173    return hasher.hexdigest()
174
175
176def expected_hash(rosetta=False):
177    """Pulls expected hash from digests file."""
178
179    expected_plat = platform_arch_normalized(rosetta)
180
181    with open(DIGESTS_FILE, 'r') as ins:
182        for line in ins:
183            line = line.strip()
184            if line.startswith('#') or not line:
185                continue
186            plat, hashtype, hashval = line.split()
187            if hashtype == 'sha256' and plat == expected_plat:
188                return hashval
189    raise Exception('platform {} not in {}'.format(expected_plat, DIGESTS_FILE))
190
191
192def https_connect_with_proxy(target_url):
193    """Create HTTPSConnection with proxy support."""
194
195    proxy_env = os.environ.get('HTTPS_PROXY') or os.environ.get('https_proxy')
196    if proxy_env in (None, ''):
197        conn = httplib.HTTPSConnection(target_url)
198        return conn
199
200    url = urlparse.urlparse(proxy_env)
201    conn = httplib.HTTPSConnection(url.hostname, url.port)
202    headers = {}
203    if url.username and url.password:
204        auth = '%s:%s' % (url.username, url.password)
205        py_version = sys.version_info.major
206        if py_version >= 3:
207            headers['Proxy-Authorization'] = 'Basic ' + str(
208                base64.b64encode(auth.encode()).decode()
209            )
210        else:
211            headers['Proxy-Authorization'] = 'Basic ' + base64.b64encode(auth)
212    conn.set_tunnel(target_url, 443, headers)
213    return conn
214
215
216def client_bytes(rosetta=False):
217    """Pull down the CIPD client and return it as a bytes object.
218
219    Often CIPD_HOST returns a 302 FOUND with a pointer to
220    storage.googleapis.com, so this needs to handle redirects, but it
221    shouldn't require the initial response to be a redirect either.
222    """
223
224    with open(VERSION_FILE, 'r') as ins:
225        version = ins.read().strip()
226
227    try:
228        conn = https_connect_with_proxy(CIPD_HOST)
229    except AttributeError:
230        print('=' * 70)
231        print(
232            '''
233It looks like this version of Python does not support SSL. This is common
234when using Homebrew. If using Homebrew please run the following commands.
235If not using Homebrew check how your version of Python was built.
236
237brew install openssl  # Probably already installed, but good to confirm.
238brew uninstall python && brew install python
239'''.strip()
240        )
241        print('=' * 70)
242        raise
243
244    full_platform = platform_arch_normalized(rosetta)
245    if full_platform not in SUPPORTED_PLATFORMS:
246        raise UnsupportedPlatform(full_platform)
247
248    path = '/client?platform={}&version={}'.format(full_platform, version)
249
250    for _ in range(10):
251        try:
252            conn.request('GET', path)
253            res = conn.getresponse()
254            # Have to read the response before making a new request, so make
255            # sure we always read it.
256            content = res.read()
257        except ssl.SSLError:
258            print(
259                '\n'
260                'Bootstrap: SSL error in Python when downloading CIPD client.\n'
261                'If using system Python try\n'
262                '\n'
263                '    sudo pip install certifi\n'
264                '\n'
265                'And if on the system Python on a Mac try\n'
266                '\n'
267                '    /Applications/Python 3.6/Install Certificates.command\n'
268                '\n'
269                'If using Homebrew Python try\n'
270                '\n'
271                '    brew install openssl\n'
272                '    brew uninstall python\n'
273                '    brew install python\n'
274                '\n'
275                "If those don't work, address all the potential issues shown \n"
276                'by the following command.\n'
277                '\n'
278                '    brew doctor\n'
279                '\n'
280                "Otherwise, check that your machine's Python can use SSL, "
281                'testing with the httplib module on Python 2 or http.client on '
282                'Python 3.',
283                file=sys.stderr,
284            )
285            raise
286
287        # Found client bytes.
288        if res.status == httplib.OK:  # pylint: disable=no-else-return
289            return content
290
291        # Redirecting to another location.
292        elif res.status == httplib.FOUND:
293            location = res.getheader('location')
294            url = urlparse.urlparse(location)
295            if url.netloc != conn.host:
296                conn = https_connect_with_proxy(url.netloc)
297            path = '{}?{}'.format(url.path, url.query)
298
299        # Some kind of error in this response.
300        else:
301            break
302
303    raise Exception(
304        'failed to download client from https://{}{}'.format(CIPD_HOST, path)
305    )
306
307
308def bootstrap(
309    client,
310    silent=('PW_ENVSETUP_QUIET' in os.environ),
311    rosetta=False,
312):
313    """Bootstrap cipd client installation."""
314
315    client_dir = os.path.dirname(client)
316    if not os.path.isdir(client_dir):
317        os.makedirs(client_dir)
318
319    if not silent:
320        print(
321            'Bootstrapping cipd client for {}'.format(
322                platform_arch_normalized(rosetta)
323            )
324        )
325
326    tmp_path = client + '.tmp'
327    with open(tmp_path, 'wb') as tmp:
328        tmp.write(client_bytes(rosetta))
329
330    expected = expected_hash(rosetta=rosetta)
331    actual = actual_hash(tmp_path)
332
333    if expected != actual:
334        raise Exception(
335            'digest of downloaded CIPD client is incorrect, '
336            'check that digests file is current'
337        )
338
339    os.chmod(tmp_path, 0o755)
340    os.rename(tmp_path, client)
341
342
343def selfupdate(client):
344    """Update cipd client."""
345
346    cmd = [
347        client,
348        'selfupdate',
349        '-version-file',
350        VERSION_FILE,
351        '-service-url',
352        'https://{}'.format(CIPD_HOST),
353    ]
354    subprocess.check_call(cmd)
355
356
357def _default_client(install_dir):
358    client = os.path.join(install_dir, 'cipd')
359    if os.name == 'nt':
360        client += '.exe'
361    return client
362
363
364def init(
365    install_dir=DEFAULT_INSTALL_DIR,
366    silent=False,
367    client=None,
368    rosetta=False,
369):
370    """Install/update cipd client."""
371
372    if not client:
373        client = _default_client(install_dir)
374
375    os.environ['CIPD_HTTP_USER_AGENT_PREFIX'] = user_agent()
376
377    if not os.path.isfile(client):
378        bootstrap(client, silent, rosetta=rosetta)
379
380    try:
381        selfupdate(client)
382    except subprocess.CalledProcessError:
383        print(
384            'CIPD selfupdate failed. Bootstrapping then retrying...',
385            file=sys.stderr,
386        )
387        bootstrap(client, rosetta=rosetta)
388        selfupdate(client)
389
390    return client
391
392
393def main(install_dir=DEFAULT_INSTALL_DIR, silent=False):
394    """Install/update cipd client."""
395
396    client = _default_client(install_dir)
397
398    try:
399        init(install_dir=install_dir, silent=silent, client=client)
400
401    except UnsupportedPlatform:
402        # Don't show help message below for this exception.
403        raise
404
405    except Exception:
406        print(
407            'Failed to initialize CIPD. Run '
408            '`CIPD_HTTP_USER_AGENT_PREFIX={user_agent}/manual {client} '
409            "selfupdate -version-file '{version_file}'` "
410            'to diagnose if this is persistent.'.format(
411                user_agent=user_agent(),
412                client=client,
413                version_file=VERSION_FILE,
414            ),
415            file=sys.stderr,
416        )
417        raise
418
419    return client
420
421
422if __name__ == '__main__':
423    client_exe = main()
424    subprocess.check_call([client_exe] + sys.argv[1:])
425