1#!/usr/bin/env python3
2# -*- coding: utf-8 -*-
3# Copyright (c) 2022-2023 Huawei Device Co., Ltd.
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#     http://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,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15
16import os
17import sys
18import argparse
19import subprocess
20import tarfile
21import zipfile
22import ssl
23import shutil
24from multiprocessing import cpu_count
25from concurrent.futures import ThreadPoolExecutor, as_completed
26from functools import partial
27from urllib.request import urlopen
28import urllib.error
29from rich.progress import (
30    BarColumn,
31    DownloadColumn,
32    Progress,
33    TaskID,
34    TextColumn,
35    TimeRemainingColumn,
36    TransferSpeedColumn,
37)
38from util import read_json_file
39
40progress = Progress(
41    TextColumn("[bold blue]{task.fields[filename]}", justify="right"),
42    BarColumn(bar_width=None),
43    "[progress.percentage]{task.percentage:>3.1f}%",
44    "•",
45    DownloadColumn(),
46    "•",
47    TransferSpeedColumn(),
48    "•",
49    TimeRemainingColumn(),
50)
51
52def _run_cmd(cmd):
53    res = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE,
54                           stderr=subprocess.PIPE)
55    sout, serr = res.communicate()
56    return sout.rstrip().decode('utf-8'), serr, res.returncode
57
58def _check_sha256(check_url, local_file):
59    check_sha256_cmd = ''.join(['curl -s -k ', check_url, '.sha256'])
60    local_sha256_cmd = ''.join(['sha256sum ', local_file, "|cut -d ' ' -f1"])
61    check_sha256, err, returncode = _run_cmd(check_sha256_cmd)
62    local_sha256, err, returncode = _run_cmd(local_sha256_cmd)
63    return check_sha256 == local_sha256
64
65def _check_sha256_by_mark(args, check_url, code_dir, unzip_dir, unzip_filename):
66    check_sha256_cmd = ''.join(['curl -s -k ', check_url, '.sha256'])
67    check_sha256, err, returncode = _run_cmd(check_sha256_cmd)
68    mark_file_dir = os.path.join(code_dir, unzip_dir)
69    mark_file_name = ''.join([check_sha256, '.', unzip_filename, '.mark'])
70    mark_file_path = os.path.join(mark_file_dir, mark_file_name)
71    args.mark_file_path = mark_file_path
72    return os.path.exists(mark_file_path)
73
74def _config_parse(config, tool_repo):
75    unzip_dir = config.get('unzip_dir')
76    huaweicloud_url = ''.join([tool_repo, config.get('file_path')])
77    unzip_filename = config.get('unzip_filename')
78    md5_huaweicloud_url_cmd = ''.join(['echo ', huaweicloud_url, "|md5sum|cut -d ' ' -f1"])
79    md5_huaweicloud_url, err, returncode = _run_cmd(md5_huaweicloud_url_cmd)
80    bin_file = os.path.basename(huaweicloud_url)
81    return unzip_dir, huaweicloud_url, unzip_filename, md5_huaweicloud_url, bin_file
82
83def _uncompress(args, src_file, code_dir, unzip_dir, unzip_filename, mark_file_path):
84    dest_dir = os.path.join(code_dir, unzip_dir)
85    if src_file[-3:] == 'zip':
86        cmd = 'unzip -o {} -d {};echo 0 > {}'.format(src_file, dest_dir, mark_file_path)
87    elif src_file[-6:] == 'tar.gz':
88        cmd = 'tar -xvzf {} -C {};echo 0 > {}'.format(src_file, dest_dir, mark_file_path)
89    else:
90        cmd = 'tar -xvf {} -C {};echo 0 > {}'.format(src_file, dest_dir, mark_file_path)
91    _, _, returncode = _run_cmd(cmd)
92    return returncode
93
94def _copy_url(args, task_id, url, local_file, code_dir, unzip_dir, unzip_filename, mark_file_path):
95    # download files
96    download_buffer_size = 32768
97    progress.console.log('Requesting {}'.format(url))
98    flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL
99    modes = 0o777
100    try:
101        response = urlopen(url)
102    except urllib.error.HTTPError as e:
103        progress.console.log("Failed to open {}, HTTPError: {}".format(url, e.code), style='red')
104        return 1
105    progress.update(task_id, total=int(response.info()["Content-length"]))
106    with os.fdopen(os.open(local_file, flags, modes), 'wb') as dest_file:
107        progress.start_task(task_id)
108        for data in iter(partial(response.read, download_buffer_size), b""):
109            dest_file.write(data)
110            progress.update(task_id, advance=len(data))
111    progress.console.log("Downloaded {}".format(local_file))
112    if not _check_sha256(url, local_file):
113        progress.console.log('{}, Sha256 check download FAILED.'.format(local_file), style='red')
114        return 1
115
116    # decompressing files
117    progress.console.log("Decompressing {}".format(local_file))
118    returncode = _uncompress(args, local_file, code_dir, unzip_dir, unzip_filename, mark_file_path)
119    progress.console.log("Decompressed {}".format(local_file))
120    return returncode
121
122
123def _hwcloud_download_wrapper(args, config, bin_dir, code_dir, retries):
124    attempt = 0
125    success = False
126    with progress:
127        while not success and attempt < retries:
128            success = _hwcloud_download(args, config, bin_dir, code_dir, retries)
129            attempt += 1
130    return success
131
132
133def _hwcloud_download(args, config, bin_dir, code_dir, retries):
134    try:
135        cnt = cpu_count()
136    except:
137        cnt = 1
138
139    success = False
140    with ThreadPoolExecutor(max_workers=cnt) as pool:
141        tasks = dict()
142        for config_info in config:
143            unzip_dir, huaweicloud_url, unzip_filename, md5_huaweicloud_url, bin_file = _config_parse(config_info,
144                args.tool_repo)
145            abs_unzip_dir = os.path.join(code_dir, unzip_dir)
146            if not os.path.exists(abs_unzip_dir):
147                os.makedirs(abs_unzip_dir)
148            if _check_sha256_by_mark(args, huaweicloud_url, code_dir, unzip_dir, unzip_filename):
149                progress.console.log('{}, Sha256 markword check OK.'.format(huaweicloud_url), style='green')
150                continue
151
152            _run_cmd(''.join(['rm -rf ', code_dir, '/', unzip_dir, '/*.', unzip_filename, '.mark']))
153            _run_cmd(''.join(['rm -rf ', code_dir, '/', unzip_dir, '/', unzip_filename]))
154            local_file = os.path.join(bin_dir, ''.join([md5_huaweicloud_url, '.', bin_file]))
155
156            if os.path.exists(local_file) and not _check_sha256(huaweicloud_url, local_file):
157                os.remove(local_file)
158
159            if not os.path.exists(local_file):
160                filename = huaweicloud_url.split("/")[-1]
161                task_id = progress.add_task("download", filename=filename, start=False)
162                task = pool.submit(_copy_url, args, task_id, huaweicloud_url, local_file, code_dir, unzip_dir,
163                    unzip_filename, args.mark_file_path)
164                tasks[task] = os.path.basename(huaweicloud_url)
165                continue
166
167            if _check_sha256(huaweicloud_url, local_file):
168                progress.console.log('{}, Sha256 check download OK.'.format(local_file), style='green')
169                task = pool.submit(_uncompress, args, local_file, code_dir, unzip_dir, unzip_filename,
170                    args.mark_file_path)
171                tasks[task] = os.path.basename(huaweicloud_url)
172            else:
173                os.remove(local_file)
174        returncode = 0
175        for task in as_completed(tasks):
176            if task.result():
177                returncode += task.result()
178            progress.console.log('{}, download and decompress completed, exit code: {}'
179                                     .format(tasks.get(task), task.result()), style='green')
180        success = returncode == 0
181    return success
182
183
184def _file_handle(config, code_dir):
185    for config_info in config:
186        src_dir = ''.join([code_dir, config_info.get('src')])
187        dest_dir = ''.join([code_dir, config_info.get('dest')])
188        tmp_dir = config_info.get('tmp')
189        symlink_src = config_info.get('symlink_src')
190        symlink_dest = config_info.get('symlink_dest')
191        if os.path.exists(src_dir):
192            if tmp_dir:
193                tmp_dir = ''.join([code_dir, tmp_dir])
194                shutil.move(src_dir, tmp_dir)
195                cmd = 'mv {}/*.mark {}'.format(dest_dir, tmp_dir)
196                _run_cmd(cmd)
197                if os.path.exists(dest_dir):
198                    shutil.rmtree(dest_dir)
199                shutil.move(tmp_dir, dest_dir)
200            elif symlink_src and symlink_dest:
201                if os.path.exists(dest_dir) and dest_dir != src_dir:
202                    shutil.rmtree(dest_dir)
203                shutil.move(src_dir, dest_dir)
204                os.symlink(''.join([dest_dir, symlink_src]), ''.join([dest_dir, symlink_dest]))
205            else:
206                _run_cmd('chmod 755 {} -R'.format(dest_dir))
207
208def main():
209    parser = argparse.ArgumentParser()
210    parser.add_argument('--skip-ssl', action='store_true', help='skip ssl authentication')
211    parser.add_argument('--tool-repo', default='https://repo.huaweicloud.com', help='prebuilt file download source')
212    parser.add_argument('--host-cpu', help='host cpu', required=True)
213    parser.add_argument('--host-platform', help='host platform', required=True)
214    args = parser.parse_args()
215    args.code_dir = os.path.abspath(os.path.join(os.getcwd()))
216    if args.skip_ssl:
217        ssl._create_default_https_context = ssl._create_unverified_context
218
219    host_platform = args.host_platform
220    host_cpu = args.host_cpu
221    tool_repo = args.tool_repo
222    config_file = os.path.join(args.code_dir,
223        'arkcompiler/toolchain/build/prebuilts_download/prebuilts_download_config.json')
224    config_info = read_json_file(config_file)
225    file_handle_config = config_info.get('file_handle_config')
226
227    args.bin_dir = os.path.join(args.code_dir, config_info.get('prebuilts_download_dir'))
228    if not os.path.exists(args.bin_dir):
229        os.makedirs(args.bin_dir)
230    copy_config = config_info.get(host_platform).get(host_cpu).get('copy_config')
231    if host_platform == 'linux':
232        linux_copy_config = config_info.get(host_platform).get(host_cpu).get('linux_copy_config')
233        copy_config.extend(linux_copy_config)
234    elif host_platform == 'darwin':
235        darwin_copy_config = config_info.get(host_platform).get(host_cpu).get('darwin_copy_config')
236        copy_config.extend(darwin_copy_config)
237    retries = config_info.get('retries')
238    args.retries = 1 if retries is None else retries
239    if not _hwcloud_download_wrapper(args, copy_config, args.bin_dir, args.code_dir, args.retries):
240        return 1
241    _file_handle(file_handle_config, args.code_dir)
242    return 0
243
244
245if __name__ == '__main__':
246    sys.exit(main())
247