third_party/pigweed/repo/pw_env_setup/py/pw_env_setup/cipd_setup/wrapper.py

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