aboutsummaryrefslogtreecommitdiffstats
path: root/ya
blob: 813ae7119b047bd054ccd04641c3e18489303f7d (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
#!/usr/bin/env python

# Please, keep this script in sync with arcadia/ya

import os
import sys
import platform

RETRIES = 5
HASH_PREFIX = 10

REGISTRY_ENDPOINT = os.environ.get("YA_REGISTRY_ENDPOINT", "https://s3.mds.yandex.net/devtools-registry")


PLATFORM_MAP = {
    "data": {
        "darwin": {
            "md5": "1356748b6bb7ce19359bc1c87d81cf5e",
            "urls": [
                f"{REGISTRY_ENDPOINT}/4948880410",
            ],
        },
        "darwin-arm64": {
            "md5": "2ba999e77660102f78c2584931eef3a2",
            "urls": [
                f"{REGISTRY_ENDPOINT}/4948879923",
            ],
        },
        "linux-aarch64": {
            "md5": "731df682969dddc89bd7f3160960ffa1",
            "urls": [
                f"{REGISTRY_ENDPOINT}/4948879460",
            ],
        },
        "win32-clang-cl": {
            "md5": "e2289cac60bf800470d858af48fce490",
            "urls": [
                f"{REGISTRY_ENDPOINT}/4948880870",
            ],
        },
        "linux": {
            "md5": "c71678ed371ae966d7c0235352a398e3",
            "urls": [
                f"{REGISTRY_ENDPOINT}/4948881174",
            ],
        },
    }
}
# End of mapping


def create_dirs(path):
    try:
        os.makedirs(path)
    except OSError as e:
        import errno

        if e.errno != errno.EEXIST:
            raise

    return path


def home_dir():
    # Do not trust $HOME, as it is unreliable in certain environments
    # Temporarily delete os.environ["HOME"] to force reading current home directory from /etc/passwd
    home_from_env = os.environ.pop("HOME", None)
    try:
        home_from_passwd = os.path.expanduser("~")
        if os.path.isabs(home_from_passwd):
            # This home dir is valid, prefer it over $HOME
            return home_from_passwd
        else:
            # When python is built with musl (this is quire weird though),
            # only users from /etc/passwd will be properly resolved,
            # as musl does not have nss module for LDAP integration.
            return home_from_env

    finally:
        if home_from_env is not None:
            os.environ["HOME"] = home_from_env


def misc_root():
    return create_dirs(os.getenv('YA_CACHE_DIR') or os.path.join(home_dir(), '.ya'))


def tool_root():
    return create_dirs(os.getenv('YA_CACHE_DIR_TOOLS') or os.path.join(misc_root(), 'tools'))


# TODO: remove when switched to S3, won't be needed in OSS
def ya_token():
    def get_token_from_file():
        try:
            with open(os.environ.get('YA_TOKEN_PATH', os.path.join(home_dir(), '.ya_token')), 'r') as f:
                return f.read().strip()
        except:
            pass

    return os.getenv('YA_TOKEN') or get_token_from_file()


TOOLS_DIR = tool_root()


def uniq(size=6):
    import string
    import random

    return ''.join(random.choice(string.ascii_lowercase + string.digits) for _ in range(size))


_ssl_is_tuned = False


def _tune_ssl():
    global _ssl_is_tuned
    if _ssl_is_tuned:
        return

    try:
        import ssl

        ssl._create_default_https_context = ssl._create_unverified_context
    except AttributeError:
        pass

    try:
        import urllib3

        urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
    except (AttributeError, ImportError):
        pass
    _ssl_is_tuned = True


def _fetch(url, into):
    import hashlib

    _tune_ssl()

    from urllib.request import urlopen
    from urllib.request import Request
    from urllib.parse import urlparse

    request = Request(str(url))
    # TODO: Remove when switched to S3 distribution
    request.add_header('User-Agent', 'ya-bootstrap')
    token = ya_token()
    if token:
        request.add_header('Authorization', 'OAuth {}'.format(token))

    md5 = hashlib.md5()
    sys.stderr.write('Downloading %s ' % url)
    sys.stderr.flush()
    conn = urlopen(request, timeout=10)
    sys.stderr.write('[')
    sys.stderr.flush()
    try:
        with open(into, 'wb') as f:
            while True:
                block = conn.read(1024 * 1024)
                sys.stderr.write('.')
                sys.stderr.flush()
                if block:
                    md5.update(block)
                    f.write(block)
                else:
                    break
        return md5.hexdigest()

    finally:
        sys.stderr.write('] ')
        sys.stderr.flush()


def _atomic_fetch(url, into, md5):
    tmp_dest = into + '.' + uniq()
    try:
        real_md5 = _fetch(url, tmp_dest)
        if real_md5 != md5:
            raise Exception('MD5 mismatched: %s differs from %s' % (real_md5, md5))
        os.rename(tmp_dest, into)
        sys.stderr.write('OK\n')
    except Exception as e:
        sys.stderr.write('ERROR: ' + str(e) + '\n')
        raise
    finally:
        try:
            os.remove(tmp_dest)
        except OSError:
            pass


def _extract(path, into):
    import tarfile

    tar = tarfile.open(path, errorlevel=2)

    # tar.extractall() will try to set file ownership according to the attributes stored in the archive
    # by calling TarFile.chown() method.
    # As this information is hardly relevant to the point of deployment / extraction,
    # it will just fail (python2) if ya is executed with root euid, or silently set non-existent numeric owner (python3)
    # to the files being extracted.
    # mock it with noop to retain current user ownership.
    tar.chown = lambda *args, **kwargs: None

    tar.extractall(path=into)
    tar.close()


def _get(urls, md5):
    dest_path = os.path.join(TOOLS_DIR, md5[:HASH_PREFIX])

    if not os.path.exists(dest_path):
        for iter in range(RETRIES):
            try:
                _atomic_fetch(urls[iter % len(urls)], dest_path, md5)
                break
            except Exception:
                if iter + 1 == RETRIES:
                    raise
                else:
                    import time

                    time.sleep(iter)

    return dest_path


def _get_dir(urls, md5, ya_name):
    dest_dir = os.path.join(TOOLS_DIR, md5[:HASH_PREFIX] + '_d')

    if os.path.isfile(os.path.join(dest_dir, ya_name)):
        return dest_dir

    try:
        packed_path = _get(urls, md5)
    except Exception:
        if os.path.isfile(os.path.join(dest_dir, ya_name)):
            return dest_dir
        raise

    tmp_dir = dest_dir + '.' + uniq()
    try:
        try:
            _extract(packed_path, tmp_dir)
        except Exception:
            if os.path.isfile(os.path.join(dest_dir, ya_name)):
                return dest_dir
            raise

        try:
            os.rename(tmp_dir, dest_dir)
        except OSError as e:
            import errno

            if e.errno != errno.ENOTEMPTY:
                raise

        return dest_dir
    finally:
        import shutil

        shutil.rmtree(tmp_dir, ignore_errors=True)
        try:
            os.remove(packed_path)
        except Exception:
            pass


def _mine_repo_root():
    # We think that this script is located in the root of the repo.
    return os.path.dirname(os.path.realpath(__file__))


def main():
    if not os.path.exists(TOOLS_DIR):
        os.makedirs(TOOLS_DIR)

    result_args = sys.argv[1:]

    meta = PLATFORM_MAP['data']
    my_platform = platform.system().lower()
    my_machine = platform.machine().lower()
    if my_platform == 'linux':
        if 'ppc64le' in platform.platform():
            my_platform = 'linux-ppc64le'
        elif 'aarch64' in platform.platform():
            my_platform = 'linux-aarch64'
        else:
            my_platform = 'linux_musl'
    if my_platform == 'darwin' and my_machine == 'arm64':
        my_platform = 'darwin-arm64'

    def _platform_key(target_platform):
        """match by max prefix length, prefer shortest"""

        def _key_for_platform(platform):
            return len(os.path.commonprefix([target_platform, platform])), -len(platform)

        return _key_for_platform

    best_key = max(meta.keys(), key=_platform_key(my_platform))
    value = meta[best_key]

    ya_name = {'win32': 'ya-bin.exe', 'win32-clang-cl': 'ya-bin.exe'}.get(best_key, 'ya-bin')  # XXX
    ya_dir = _get_dir(value['urls'], value['md5'], ya_name)

    # Popen `args` must have `str` type
    ya_path = str(os.path.join(ya_dir, ya_name))

    env = os.environ.copy()
    if 'YA_SOURCE_ROOT' not in env:
        src_root = _mine_repo_root()
        if src_root is not None:
            env['YA_SOURCE_ROOT'] = src_root

    for env_name in [
        'LD_PRELOAD',
        'Y_PYTHON_SOURCE_ROOT',
    ]:
        if env_name in os.environ:
            sys.stderr.write(
                "Warn: {}='{}' is specified and may affect the correct operation of the ya\n".format(
                    env_name, env[env_name]
                )
            )

    if os.name == 'nt':
        import subprocess

        p = subprocess.Popen([ya_path] + result_args, env=env)
        p.wait()
        sys.exit(p.returncode)
    else:
        os.execve(ya_path, [ya_path] + result_args, env)


if __name__ == '__main__':
    try:
        main()
    except Exception as e:
        sys.stderr.write('ERROR: ' + str(e) + '\n')
        from traceback import format_exc

        sys.stderr.write(format_exc() + "\n")
        sys.exit(1)