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