| #!/usr/bin/env python |
| # Copyright 2020 The Cobalt Authors. All Rights Reserved. |
| # |
| # Licensed under the Apache License, Version 2.0 (the "License"); |
| # you may not use this file except in compliance with the License. |
| # You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| |
| import argparse |
| import hashlib |
| import logging |
| import os |
| import shutil |
| import ssl |
| import stat |
| import tempfile |
| try: |
| import urllib.request as urllib |
| except ImportError: |
| import urllib2 as urllib |
| |
| # ssl.SSLContext (and thus ssl.create_default_context) was introduced in |
| # python 2.7.9. depot_tools provides python 2.7.6, so we wrap this import. |
| try: |
| from ssl import create_default_context |
| except ImportError: |
| create_default_context = lambda: None |
| |
| _BASE_GCS_URL = 'https://storage.googleapis.com' |
| _BUFFER_SIZE = 2 * 1024 * 1024 |
| |
| |
| def AddExecutableBits(filename): |
| st = os.stat(filename) |
| os.chmod(filename, st.st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH) |
| |
| |
| def ExtractSha1(filename): |
| with open(filename, 'rb') as f: |
| sha1 = hashlib.sha1() |
| buf = f.read(_BUFFER_SIZE) |
| while buf: |
| sha1.update(buf) |
| buf = f.read(_BUFFER_SIZE) |
| return sha1.hexdigest() |
| |
| |
| def _DownloadFromGcsAndCheckSha1(bucket, sha1): |
| url = '{}/{}/{}'.format(_BASE_GCS_URL, bucket, sha1) |
| context = create_default_context() |
| |
| res = urllib.urlopen(url, context=context) if context else urllib.urlopen(url) |
| if not res: |
| logging.error('Could not reach %s', url) |
| return None |
| |
| with tempfile.NamedTemporaryFile(delete=False) as tmp_file: |
| shutil.copyfileobj(res, tmp_file) |
| |
| if ExtractSha1(tmp_file.name) != sha1: |
| logging.error('Local and remote sha1s do not match. Skipping download.') |
| return None |
| |
| return tmp_file |
| |
| |
| def MaybeDownloadFileFromGcs(bucket, sha1_file, output_file, force=False): |
| """Download file from GCS if it doesn't already exist. |
| |
| Args: |
| bucket: The GCS bucket that the file is stored in. |
| sha1_file: The file containing the sha1 of the file to download |
| output_file: The file to output to |
| force: If true, will overwrite an existing output_file |
| Returns: |
| True if it writes a file, False otherwise. |
| """ |
| if not os.path.exists(sha1_file): |
| logging.error("Provided sha1 file %s doesn't exist", sha1_file) |
| with open(sha1_file) as fd: |
| sha1 = fd.read().strip() |
| |
| if not force and os.path.exists(output_file): |
| with open(output_file, 'rb') as f: |
| if hashlib.sha1(f.read()).hexdigest() == sha1: |
| logging.info('%s exists and sha1s match, skipping download', |
| output_file) |
| return False |
| |
| tmp_file = _DownloadFromGcsAndCheckSha1(bucket, sha1) |
| if not tmp_file: |
| return False |
| |
| shutil.move(tmp_file.name, output_file) |
| AddExecutableBits(output_file) |
| return True |
| |
| |
| def MaybeDownloadDirectoryFromGcs(bucket, |
| sha1_directory, |
| output_directory, |
| force=False): |
| res = True |
| for filename in os.listdir(sha1_directory): |
| filebase, ext = os.path.splitext(filename) |
| if ext == '.sha1': |
| filepath = os.path.join(sha1_directory, filename) |
| output_filepath = os.path.join(output_directory, filebase) |
| if not MaybeDownloadFileFromGcs(bucket, filepath, output_filepath, force): |
| res = False |
| return res |
| |
| |
| if __name__ == "__main__": |
| logging_format = '[%(levelname)s:%(filename)s:%(lineno)s] %(message)s' |
| logging.basicConfig( |
| level=logging.INFO, format=logging_format, datefmt='%H:%M:%S') |
| |
| parser = argparse.ArgumentParser() |
| parser.add_argument( |
| '-b', |
| '--bucket', |
| required=True, |
| help='GCS bucket the resource is located in.') |
| parser.add_argument( |
| '-s', |
| '--sha1', |
| required=True, |
| help='Path to file or directory containing sha1 checksum(s).') |
| parser.add_argument( |
| '-o', |
| '--output', |
| required=True, |
| help='Path to file or directory to download file(s) to.') |
| parser.add_argument( |
| '-f', |
| '--force', |
| action='store_true', |
| help='Replace an existing resource.') |
| args = parser.parse_args() |
| |
| if os.path.isdir(args.sha1): |
| MaybeDownloadDirectoryFromGcs(args.bucket, args.sha1, args.output, |
| args.force) |
| else: |
| MaybeDownloadFileFromGcs(args.bucket, args.sha1, args.output, args.force) |