You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
145 lines
4.7 KiB
Python
145 lines
4.7 KiB
Python
#!/usr/bin/env python3
|
|
# Copyright (c) 2024 The Chromium Authors. All rights reserved.
|
|
# Use of this source code is governed by a BSD-style license that can be
|
|
# found in the LICENSE file.
|
|
"""Download files from Google Storage, given the bucket and file."""
|
|
|
|
import optparse
|
|
import os
|
|
import re
|
|
import sys
|
|
|
|
import subprocess2
|
|
|
|
# Env vars that tempdir can be gotten from; minimally, this
|
|
# needs to match python's tempfile module and match normal
|
|
# unix standards.
|
|
_TEMPDIR_ENV_VARS = ('TMPDIR', 'TEMP', 'TMP')
|
|
|
|
GSUTIL_DEFAULT_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)),
|
|
'gsutil.py')
|
|
|
|
# Maps sys.platform to what we actually want to call them.
|
|
PLATFORM_MAPPING = {
|
|
'cygwin': 'win',
|
|
'darwin': 'mac',
|
|
'linux': 'linux', # Python 3.3+.
|
|
'win32': 'win',
|
|
'aix6': 'aix',
|
|
'aix7': 'aix',
|
|
'zos': 'zos',
|
|
}
|
|
|
|
|
|
def GetNormalizedPlatform():
|
|
"""Returns the result of sys.platform accounting for cygwin.
|
|
Under cygwin, this will always return "win32" like the native Python."""
|
|
if sys.platform == 'cygwin':
|
|
return 'win32'
|
|
return sys.platform
|
|
|
|
|
|
# Common utilities
|
|
class Gsutil(object):
|
|
"""Call gsutil with some predefined settings. This is a convenience object,
|
|
and is also immutable.
|
|
|
|
HACK: This object is used directly by the external script
|
|
`<depot_tools>/win_toolchain/get_toolchain_if_necessary.py`
|
|
"""
|
|
|
|
MAX_TRIES = 5
|
|
RETRY_BASE_DELAY = 5.0
|
|
RETRY_DELAY_MULTIPLE = 1.3
|
|
VPYTHON3 = ('vpython3.bat'
|
|
if GetNormalizedPlatform() == 'win32' else 'vpython3')
|
|
|
|
def __init__(self, path, boto_path=None):
|
|
if not os.path.exists(path):
|
|
raise FileNotFoundError('GSUtil not found in %s' % path)
|
|
self.path = path
|
|
self.boto_path = boto_path
|
|
|
|
def get_sub_env(self):
|
|
env = os.environ.copy()
|
|
if self.boto_path == os.devnull:
|
|
env['AWS_CREDENTIAL_FILE'] = ''
|
|
env['BOTO_CONFIG'] = ''
|
|
elif self.boto_path:
|
|
env['AWS_CREDENTIAL_FILE'] = self.boto_path
|
|
env['BOTO_CONFIG'] = self.boto_path
|
|
|
|
if PLATFORM_MAPPING[sys.platform] != 'win':
|
|
env.update((x, "/tmp") for x in _TEMPDIR_ENV_VARS)
|
|
|
|
return env
|
|
|
|
def call(self, *args):
|
|
cmd = [self.VPYTHON3, self.path]
|
|
cmd.extend(args)
|
|
return subprocess2.call(cmd, env=self.get_sub_env())
|
|
|
|
def check_call(self, *args):
|
|
cmd = [self.VPYTHON3, self.path]
|
|
cmd.extend(args)
|
|
((out, err), code) = subprocess2.communicate(cmd,
|
|
stdout=subprocess2.PIPE,
|
|
stderr=subprocess2.PIPE,
|
|
env=self.get_sub_env())
|
|
|
|
out = out.decode('utf-8', 'replace')
|
|
err = err.decode('utf-8', 'replace')
|
|
|
|
# Parse output.
|
|
status_code_match = re.search('status=([0-9]+)', err)
|
|
if status_code_match:
|
|
return (int(status_code_match.group(1)), out, err)
|
|
if ('ServiceException: 401 Anonymous' in err):
|
|
return (401, out, err)
|
|
if ('You are attempting to access protected data with '
|
|
'no configured credentials.' in err):
|
|
return (403, out, err)
|
|
if 'matched no objects' in err or 'No URLs matched' in err:
|
|
return (404, out, err)
|
|
return (code, out, err)
|
|
|
|
def check_call_with_retries(self, *args):
|
|
delay = self.RETRY_BASE_DELAY
|
|
for i in range(self.MAX_TRIES):
|
|
code, out, err = self.check_call(*args)
|
|
if not code or i == self.MAX_TRIES - 1:
|
|
break
|
|
|
|
time.sleep(delay)
|
|
delay *= self.RETRY_DELAY_MULTIPLE
|
|
|
|
return code, out, err
|
|
|
|
|
|
def main(args):
|
|
parser = optparse.OptionParser()
|
|
parser.add_option('-b',
|
|
'--bucket',
|
|
help='Google Storage bucket to fetch from.')
|
|
parser.add_option('-p', '--file', help='Path of file to fetch.')
|
|
parser.add_option('-o',
|
|
'--output',
|
|
help='Path where GCS contents should be downloaded.')
|
|
parser.add_option('-e', '--boto', help='Specify a custom boto file.')
|
|
(options, args) = parser.parse_args()
|
|
|
|
file_url = 'gs://%s/%s' % (options.bucket, options.file)
|
|
|
|
# Make sure gsutil exists where we expect it to.
|
|
if os.path.exists(GSUTIL_DEFAULT_PATH):
|
|
gsutil = Gsutil(GSUTIL_DEFAULT_PATH, boto_path=options.boto)
|
|
else:
|
|
parser.error('gsutil not found in %s, bad depot_tools checkout?' %
|
|
GSUTIL_DEFAULT_PATH)
|
|
|
|
gsutil.check_call('cp', file_url, options.output)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
sys.exit(main(sys.argv))
|