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 or updates prebuilt tools.
16
17Must be tested with Python 2 and Python 3.
18
19The stdout of this script is meant to be executed by the invoking shell.
20"""
21
22from __future__ import print_function
23
24import argparse
25import json
26import os
27import re
28import shutil
29import subprocess
30import sys
31import tempfile
32
33
34def parse(argv=None):
35    """Parse arguments."""
36
37    script_root = os.path.join(os.environ['PW_ROOT'], 'pw_env_setup', 'py',
38                               'pw_env_setup', 'cipd_setup')
39    git_root = subprocess.check_output(
40        ('git', 'rev-parse', '--show-toplevel'),
41        cwd=script_root,
42    ).decode('utf-8').strip()
43
44    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
45    parser.add_argument(
46        '--install-dir',
47        dest='root_install_dir',
48        default=os.path.join(git_root, '.cipd'),
49    )
50    parser.add_argument('--package-file',
51                        dest='package_files',
52                        metavar='PACKAGE_FILE',
53                        action='append')
54    parser.add_argument('--cipd',
55                        default=os.path.join(script_root, 'wrapper.py'))
56    parser.add_argument('--cache-dir',
57                        default=os.environ.get(
58                            'CIPD_CACHE_DIR',
59                            os.path.expanduser('~/.cipd-cache-dir')))
60
61    return parser.parse_args(argv)
62
63
64def check_auth(cipd, package_files, spin):
65    """Check have access to CIPD pigweed directory."""
66
67    paths = []
68    for package_file in package_files:
69        with open(package_file, 'r') as ins:
70            # This is an expensive RPC, so only check the first few entries
71            # in each file.
72            for i, entry in enumerate(json.load(ins)):
73                if i >= 3:
74                    break
75                parts = entry['path'].split('/')
76                while '${' in parts[-1]:
77                    parts.pop(-1)
78                paths.append('/'.join(parts))
79
80    username = None
81    try:
82        output = subprocess.check_output([cipd, 'auth-info'],
83                                         stderr=subprocess.STDOUT).decode()
84        logged_in = True
85
86        match = re.search(r'Logged in as (\S*)\.', output)
87        if match:
88            username = match.group(1)
89
90    except subprocess.CalledProcessError:
91        logged_in = False
92
93    def _check_all_paths():
94        inaccessible_paths = []
95
96        for path in paths:
97            # Not catching CalledProcessError because 'cipd ls' seems to never
98            # return an error code unless it can't reach the CIPD server.
99            output = subprocess.check_output(
100                [cipd, 'ls', path], stderr=subprocess.STDOUT).decode()
101            if 'No matching packages' not in output:
102                continue
103
104            # 'cipd ls' only lists sub-packages but ignores any packages at the
105            # given path. 'cipd instances' will give versions of that package.
106            # 'cipd instances' does use an error code if there's no such package
107            # or that package is inaccessible.
108            try:
109                subprocess.check_output([cipd, 'instances', path],
110                                        stderr=subprocess.STDOUT)
111            except subprocess.CalledProcessError:
112                inaccessible_paths.append(path)
113
114        return inaccessible_paths
115
116    inaccessible_paths = _check_all_paths()
117
118    if inaccessible_paths and not logged_in:
119        with spin.pause():
120            stderr = lambda *args: print(*args, file=sys.stderr)
121            stderr()
122            stderr('No access to the following CIPD paths:')
123            for path in inaccessible_paths:
124                stderr('  {}'.format(path))
125            stderr()
126            stderr('Attempting CIPD login')
127            try:
128                subprocess.check_call([cipd, 'auth-login'])
129            except subprocess.CalledProcessError:
130                stderr('CIPD login failed')
131                return False
132
133        inaccessible_paths = _check_all_paths()
134
135    if inaccessible_paths:
136        stderr = lambda *args: print(*args, file=sys.stderr)
137        stderr('=' * 60)
138        username_part = ''
139        if username:
140            username_part = '({}) '.format(username)
141        stderr('Your account {}does not have access to the following '
142               'paths'.format(username_part))
143        for path in inaccessible_paths:
144            stderr('  {}'.format(path))
145        stderr('=' * 60)
146        return False
147
148    return True
149
150
151def write_ensure_file(package_file, ensure_file):
152    with open(package_file, 'r') as ins:
153        data = json.load(ins)
154
155    # TODO(pwbug/103) Remove 30 days after bug fixed.
156    if os.path.isdir(ensure_file):
157        shutil.rmtree(ensure_file)
158
159    with open(ensure_file, 'w') as outs:
160        outs.write('$VerifiedPlatform linux-amd64\n'
161                   '$VerifiedPlatform mac-amd64\n'
162                   '$ParanoidMode CheckPresence\n')
163
164        for entry in data:
165            outs.write('@Subdir {}\n'.format(entry.get('subdir', '')))
166            outs.write('{} {}\n'.format(entry['path'],
167                                        ' '.join(entry['tags'])))
168
169
170def update(
171    cipd,
172    package_files,
173    root_install_dir,
174    cache_dir,
175    env_vars=None,
176    spin=None,
177):
178    """Grab the tools listed in ensure_files."""
179
180    if not check_auth(cipd, package_files, spin):
181        return False
182
183    # TODO(mohrr) use os.makedirs(..., exist_ok=True).
184    if not os.path.isdir(root_install_dir):
185        os.makedirs(root_install_dir)
186
187    if env_vars:
188        env_vars.prepend('PATH', root_install_dir)
189        env_vars.set('PW_CIPD_INSTALL_DIR', root_install_dir)
190        env_vars.set('CIPD_CACHE_DIR', cache_dir)
191
192    pw_root = None
193    if env_vars:
194        pw_root = env_vars.get('PW_ROOT', None)
195    if not pw_root:
196        pw_root = os.environ['PW_ROOT']
197
198    # Run cipd for each json file.
199    for package_file in package_files:
200        if os.path.splitext(package_file)[1] == '.ensure':
201            ensure_file = package_file
202        else:
203            ensure_file = os.path.join(
204                root_install_dir,
205                os.path.basename(
206                    os.path.splitext(package_file)[0] + '.ensure'))
207            write_ensure_file(package_file, ensure_file)
208
209        install_dir = os.path.join(
210            root_install_dir,
211            os.path.basename(os.path.splitext(package_file)[0]))
212
213        cmd = [
214            cipd,
215            'ensure',
216            '-ensure-file', ensure_file,
217            '-root', install_dir,
218            '-log-level', 'warning',
219            '-cache-dir', cache_dir,
220            '-max-threads', '0',  # 0 means use CPU count.
221        ]  # yapf: disable
222
223        # TODO(pwbug/135) Use function from common utility module.
224        with tempfile.TemporaryFile(mode='w+') as temp:
225            print(*cmd, file=temp)
226            try:
227                subprocess.check_call(cmd,
228                                      stdout=temp,
229                                      stderr=subprocess.STDOUT)
230            except subprocess.CalledProcessError:
231                temp.seek(0)
232                sys.stderr.write(temp.read())
233                raise
234
235        # Set environment variables so tools can later find things under, for
236        # example, 'share'.
237        name = os.path.basename(install_dir)
238
239        if env_vars:
240            # Some executables get installed at top-level and some get
241            # installed under 'bin'.
242            env_vars.prepend('PATH', install_dir)
243            env_vars.prepend('PATH', os.path.join(install_dir, 'bin'))
244            env_vars.set('PW_{}_CIPD_INSTALL_DIR'.format(name.upper()),
245                         install_dir)
246
247            # Windows has its own special toolchain.
248            if os.name == 'nt':
249                env_vars.prepend('PATH',
250                                 os.path.join(install_dir, 'mingw64', 'bin'))
251
252    return True
253
254
255if __name__ == '__main__':
256    update(**vars(parse()))
257    sys.exit(0)
258