vendor/depot_tools/download_from_google_storage.py

   1 #!/usr/bin/env python
   2 # Copyright (c) 2012 The Chromium Authors. All rights reserved.
   3 # Use of this source code is governed by a BSD-style license that can be
   4 # found in the LICENSE file.
   5
   6 """Download files from Google Storage based on SHA1 sums."""
   7
   8
   9 import hashlib
  10 import optparse
  11 import os
  12 import Queue
  13 import re
  14 import stat
  15 import sys
  16 import threading
  17 import time
  18
  19 import subprocess2
  20
  21
  22 GSUTIL_DEFAULT_PATH = os.path.join(
  23     os.path.dirname(os.path.abspath(__file__)),
  24     'third_party', 'gsutil', 'gsutil')
  25 # Maps sys.platform to what we actually want to call them.
  26 PLATFORM_MAPPING = {
  27     'cygwin': 'win',
  28     'darwin': 'mac',
  29     'linux2': 'linux',
  30     'win32': 'win',
  31 }
  32
  33
  34 class FileNotFoundError(IOError):
  35   pass
  36
  37
  38 class InvalidFileError(IOError):
  39   pass
  40
  41
  42 class InvalidPlatformError(Exception):
  43   pass
  44
  45
  46 def GetNormalizedPlatform():
  47   """Returns the result of sys.platform accounting for cygwin.
  48   Under cygwin, this will always return "win32" like the native Python."""
  49   if sys.platform == 'cygwin':
  50     return 'win32'
  51   return sys.platform
  52
  53
  54 # Common utilities
  55 class Gsutil(object):
  56   """Call gsutil with some predefined settings.  This is a convenience object,
  57   and is also immutable."""
  58   def __init__(self, path, boto_path, timeout=None, bypass_prodaccess=False):
  59     if not os.path.exists(path):
  60       raise FileNotFoundError('GSUtil not found in %s' % path)
  61     self.path = path
  62     self.timeout = timeout
  63     self.boto_path = boto_path
  64     self.bypass_prodaccess = bypass_prodaccess
  65
  66   def get_sub_env(self):
  67     env = os.environ.copy()
  68     if self.boto_path == os.devnull:
  69       env['AWS_CREDENTIAL_FILE'] = ''
  70       env['BOTO_CONFIG'] = ''
  71     elif self.boto_path:
  72       env['AWS_CREDENTIAL_FILE'] = self.boto_path
  73       env['BOTO_CONFIG'] = self.boto_path
  74     else:
  75       custompath = env.get('AWS_CREDENTIAL_FILE', '~/.boto') + '.depot_tools'
  76       custompath = os.path.expanduser(custompath)
  77       if os.path.exists(custompath):
  78         env['AWS_CREDENTIAL_FILE'] = custompath
  79
  80     return env
  81
  82   def call(self, *args):
  83     cmd = [sys.executable, self.path]
  84     if self.bypass_prodaccess:
  85       cmd.append('--bypass_prodaccess')
  86     cmd.extend(args)
  87     return subprocess2.call(cmd, env=self.get_sub_env(), timeout=self.timeout)
  88
  89   def check_call(self, *args):
  90     cmd = [sys.executable, self.path]
  91     if self.bypass_prodaccess:
  92       cmd.append('--bypass_prodaccess')
  93     cmd.extend(args)
  94     ((out, err), code) = subprocess2.communicate(
  95         cmd,
  96         stdout=subprocess2.PIPE,
  97         stderr=subprocess2.PIPE,
  98         env=self.get_sub_env(),
  99         timeout=self.timeout)
 100
 101     # Parse output.
 102     status_code_match = re.search('status=([0-9]+)', err)
 103     if status_code_match:
 104       return (int(status_code_match.group(1)), out, err)
 105     if ('You are attempting to access protected data with '
 106           'no configured credentials.' in err):
 107       return (403, out, err)
 108     if 'No such object' in err:
 109       return (404, out, err)
 110     return (code, out, err)
 111
 112
 113 def check_bucket_permissions(base_url, gsutil):
 114   code, _, ls_err = gsutil.check_call('ls', base_url)
 115   if code != 0:
 116     print >> sys.stderr, ls_err
 117   if code == 403:
 118     print >> sys.stderr, 'Got error 403 while authenticating to %s.' % base_url
 119     print >> sys.stderr, 'Try running "download_from_google_storage --config".'
 120   elif code == 404:
 121     print >> sys.stderr, '%s not found.' % base_url
 122   return code
 123
 124
 125 def check_platform(target):
 126   """Checks if any parent directory of target matches (win|mac|linux)."""
 127   assert os.path.isabs(target)
 128   root, target_name = os.path.split(target)
 129   if not target_name:
 130     return None
 131   if target_name in ('linux', 'mac', 'win'):
 132     return target_name
 133   return check_platform(root)
 134
 135
 136 def get_sha1(filename):
 137   sha1 = hashlib.sha1()
 138   with open(filename, 'rb') as f:
 139     while True:
 140       # Read in 1mb chunks, so it doesn't all have to be loaded into memory.
 141       chunk = f.read(1024*1024)
 142       if not chunk:
 143         break
 144       sha1.update(chunk)
 145   return sha1.hexdigest()
 146
 147
 148 # Download-specific code starts here
 149
 150 def enumerate_work_queue(input_filename, work_queue, directory,
 151                          recursive, ignore_errors, output, sha1_file,
 152                          auto_platform):
 153   if sha1_file:
 154     if not os.path.exists(input_filename):
 155       if not ignore_errors:
 156         raise FileNotFoundError('%s not found.' % input_filename)
 157       print >> sys.stderr, '%s not found.' % input_filename
 158     with open(input_filename, 'rb') as f:
 159       sha1_match = re.match('^([A-Za-z0-9]{40})$', f.read(1024).rstrip())
 160       if sha1_match:
 161         work_queue.put((sha1_match.groups(1)[0], output))
 162         return 1
 163     if not ignore_errors:
 164       raise InvalidFileError('No sha1 sum found in %s.' % input_filename)
 165     print >> sys.stderr, 'No sha1 sum found in %s.' % input_filename
 166     return 0
 167
 168   if not directory:
 169     work_queue.put((input_filename, output))
 170     return 1
 171
 172   work_queue_size = 0
 173   for root, dirs, files in os.walk(input_filename):
 174     if not recursive:
 175       for item in dirs[:]:
 176         dirs.remove(item)
 177     else:
 178       for exclude in ['.svn', '.git']:
 179         if exclude in dirs:
 180           dirs.remove(exclude)
 181     for filename in files:
 182       full_path = os.path.join(root, filename)
 183       if full_path.endswith('.sha1'):
 184         if auto_platform:
 185           # Skip if the platform does not match.
 186           target_platform = check_platform(os.path.abspath(full_path))
 187           if not target_platform:
 188             err = ('--auto_platform passed in but no platform name found in '
 189                    'the path of %s' % full_path)
 190             if not ignore_errors:
 191               raise InvalidFileError(err)
 192             print >> sys.stderr, err
 193             continue
 194           current_platform = PLATFORM_MAPPING[sys.platform]
 195           if current_platform != target_platform:
 196             continue
 197
 198         with open(full_path, 'rb') as f:
 199           sha1_match = re.match('^([A-Za-z0-9]{40})$', f.read(1024).rstrip())
 200         if sha1_match:
 201           work_queue.put(
 202               (sha1_match.groups(1)[0], full_path.replace('.sha1', '')))
 203           work_queue_size += 1
 204         else:
 205           if not ignore_errors:
 206             raise InvalidFileError('No sha1 sum found in %s.' % filename)
 207           print >> sys.stderr, 'No sha1 sum found in %s.' % filename
 208   return work_queue_size
 209
 210
 211 def _downloader_worker_thread(thread_num, q, force, base_url,
 212                               gsutil, out_q, ret_codes, verbose):
 213   while True:
 214     input_sha1_sum, output_filename = q.get()
 215     if input_sha1_sum is None:
 216       return
 217     if os.path.exists(output_filename) and not force:
 218       if get_sha1(output_filename) == input_sha1_sum:
 219         if verbose:
 220           out_q.put(
 221               '%d> File %s exists and SHA1 matches. Skipping.' % (
 222                   thread_num, output_filename))
 223         continue
 224     # Check if file exists.
 225     file_url = '%s/%s' % (base_url, input_sha1_sum)
 226     if gsutil.check_call('ls', file_url)[0] != 0:
 227       out_q.put('%d> File %s for %s does not exist, skipping.' % (
 228           thread_num, file_url, output_filename))
 229       ret_codes.put((1, 'File %s for %s does not exist.' % (
 230           file_url, output_filename)))
 231       continue
 232     # Fetch the file.
 233     out_q.put('%d> Downloading %s...' % (thread_num, output_filename))
 234     try:
 235       os.remove(output_filename)  # Delete the file if it exists already.
 236     except OSError:
 237       if os.path.exists(output_filename):
 238         out_q.put('%d> Warning: deleting %s failed.' % (
 239             thread_num, output_filename))
 240     code, _, err = gsutil.check_call('cp', '-q', file_url, output_filename)
 241     if code != 0:
 242       out_q.put('%d> %s' % (thread_num, err))
 243       ret_codes.put((code, err))
 244
 245     # Set executable bit.
 246     if sys.platform == 'cygwin':
 247       # Under cygwin, mark all files as executable. The executable flag in
 248       # Google Storage will not be set when uploading from Windows, so if
 249       # this script is running under cygwin and we're downloading an
 250       # executable, it will be unrunnable from inside cygwin without this.
 251       st = os.stat(output_filename)
 252       os.chmod(output_filename, st.st_mode | stat.S_IEXEC)
 253     elif sys.platform != 'win32':
 254       # On non-Windows platforms, key off of the custom header
 255       # "x-goog-meta-executable".
 256       #
 257       # TODO(hinoka): It is supposedly faster to use "gsutil stat" but that
 258       # doesn't appear to be supported by the gsutil currently in our tree. When
 259       # we update, this code should use that instead of "gsutil ls -L".
 260       code, out, _ = gsutil.check_call('ls', '-L', file_url)
 261       if code != 0:
 262         out_q.put('%d> %s' % (thread_num, err))
 263         ret_codes.put((code, err))
 264       elif re.search('x-goog-meta-executable:', out):
 265         st = os.stat(output_filename)
 266         os.chmod(output_filename, st.st_mode | stat.S_IEXEC)
 267
 268 def printer_worker(output_queue):
 269   while True:
 270     line = output_queue.get()
 271     # Its plausible we want to print empty lines.
 272     if line is None:
 273       break
 274     print line
 275
 276
 277 def download_from_google_storage(
 278     input_filename, base_url, gsutil, num_threads, directory, recursive,
 279     force, output, ignore_errors, sha1_file, verbose, auto_platform):
 280   # Start up all the worker threads.
 281   all_threads = []
 282   download_start = time.time()
 283   stdout_queue = Queue.Queue()
 284   work_queue = Queue.Queue()
 285   ret_codes = Queue.Queue()
 286   ret_codes.put((0, None))
 287   for thread_num in range(num_threads):
 288     t = threading.Thread(
 289         target=_downloader_worker_thread,
 290         args=[thread_num, work_queue, force, base_url,
 291               gsutil, stdout_queue, ret_codes, verbose])
 292     t.daemon = True
 293     t.start()
 294     all_threads.append(t)
 295   printer_thread = threading.Thread(target=printer_worker, args=[stdout_queue])
 296   printer_thread.daemon = True
 297   printer_thread.start()
 298
 299   # Enumerate our work queue.
 300   work_queue_size = enumerate_work_queue(
 301       input_filename, work_queue, directory, recursive,
 302       ignore_errors, output, sha1_file, auto_platform)
 303   for _ in all_threads:
 304     work_queue.put((None, None))  # Used to tell worker threads to stop.
 305
 306   # Wait for all downloads to finish.
 307   for t in all_threads:
 308     t.join()
 309   stdout_queue.put(None)
 310   printer_thread.join()
 311
 312   # See if we ran into any errors.
 313   max_ret_code = 0
 314   for ret_code, message in ret_codes.queue:
 315     max_ret_code = max(ret_code, max_ret_code)
 316     if message:
 317       print >> sys.stderr, message
 318   if verbose and not max_ret_code:
 319     print 'Success!'
 320
 321   if verbose:
 322     print 'Downloading %d files took %1f second(s)' % (
 323         work_queue_size, time.time() - download_start)
 324   return max_ret_code
 325
 326
 327 def main(args):
 328   usage = ('usage: %prog [options] target\n'
 329            'Target must be:\n'
 330            '  (default) a sha1 sum ([A-Za-z0-9]{40}).\n'
 331            '  (-s or --sha1_file) a .sha1 file, containing a sha1 sum on '
 332            'the first line.\n'
 333            '  (-d or --directory) A directory to scan for .sha1 files.')
 334   parser = optparse.OptionParser(usage)
 335   parser.add_option('-o', '--output',
 336                     help='Specify the output file name. Defaults to: '
 337                          '(a) Given a SHA1 hash, the name is the SHA1 hash. '
 338                          '(b) Given a .sha1 file or directory, the name will '
 339                          'match (.*).sha1.')
 340   parser.add_option('-b', '--bucket',
 341                     help='Google Storage bucket to fetch from.')
 342   parser.add_option('-e', '--boto',
 343                     help='Specify a custom boto file.')
 344   parser.add_option('-c', '--no_resume', action='store_true',
 345                     help='Resume download if file is partially downloaded.')
 346   parser.add_option('-f', '--force', action='store_true',
 347                     help='Force download even if local file exists.')
 348   parser.add_option('-i', '--ignore_errors', action='store_true',
 349                     help='Don\'t throw error if we find an invalid .sha1 file.')
 350   parser.add_option('-r', '--recursive', action='store_true',
 351                     help='Scan folders recursively for .sha1 files. '
 352                          'Must be used with -d/--directory')
 353   parser.add_option('-t', '--num_threads', default=1, type='int',
 354                     help='Number of downloader threads to run.')
 355   parser.add_option('-d', '--directory', action='store_true',
 356                     help='The target is a directory.  '
 357                          'Cannot be used with -s/--sha1_file.')
 358   parser.add_option('-s', '--sha1_file', action='store_true',
 359                     help='The target is a file containing a sha1 sum.  '
 360                          'Cannot be used with -d/--directory.')
 361   parser.add_option('-g', '--config', action='store_true',
 362                     help='Alias for "gsutil config".  Run this if you want '
 363                          'to initialize your saved Google Storage '
 364                          'credentials.  This will create a read-only '
 365                          'credentials file in ~/.boto.depot_tools.')
 366   parser.add_option('-n', '--no_auth', action='store_true',
 367                     help='Skip auth checking.  Use if it\'s known that the '
 368                          'target bucket is a public bucket.')
 369   parser.add_option('-p', '--platform',
 370                     help='A regular expression that is compared against '
 371                          'Python\'s sys.platform. If this option is specified, '
 372                          'the download will happen only if there is a match.')
 373   parser.add_option('-a', '--auto_platform',
 374                     action='store_true',
 375                     help='Detects if any parent folder of the target matches '
 376                          '(linux|mac|win).  If so, the script will only '
 377                          'process files that are in the paths that '
 378                          'that matches the current platform.')
 379   parser.add_option('-v', '--verbose', action='store_true',
 380                     help='Output extra diagnostic and progress information.')
 381
 382   (options, args) = parser.parse_args()
 383
 384   # Make sure we should run at all based on platform matching.
 385   if options.platform:
 386     if options.auto_platform:
 387       parser.error('--platform can not be specified with --auto_platform')
 388     if not re.match(options.platform, GetNormalizedPlatform()):
 389       if options.verbose:
 390         print('The current platform doesn\'t match "%s", skipping.' %
 391               options.platform)
 392       return 0
 393
 394   # Set the boto file to /dev/null if we don't need auth.
 395   if options.no_auth:
 396     options.boto = os.devnull
 397
 398   # Make sure gsutil exists where we expect it to.
 399   if os.path.exists(GSUTIL_DEFAULT_PATH):
 400     gsutil = Gsutil(GSUTIL_DEFAULT_PATH,
 401                     boto_path=options.boto,
 402                     bypass_prodaccess=options.no_auth)
 403   else:
 404     parser.error('gsutil not found in %s, bad depot_tools checkout?' %
 405                  GSUTIL_DEFAULT_PATH)
 406
 407   # Passing in -g/--config will run our copy of GSUtil, then quit.
 408   if options.config:
 409     return gsutil.call('config', '-r', '-o',
 410                        os.path.expanduser('~/.boto.depot_tools'))
 411
 412   if not args:
 413     parser.error('Missing target.')
 414   if len(args) > 1:
 415     parser.error('Too many targets.')
 416   if not options.bucket:
 417     parser.error('Missing bucket.  Specify bucket with --bucket.')
 418   if options.sha1_file and options.directory:
 419     parser.error('Both --directory and --sha1_file are specified, '
 420                  'can only specify one.')
 421   if options.recursive and not options.directory:
 422     parser.error('--recursive specified but --directory not specified.')
 423   if options.output and options.directory:
 424     parser.error('--directory is specified, so --output has no effect.')
 425   if (not (options.sha1_file or options.directory)
 426       and options.auto_platform):
 427     parser.error('--auto_platform must be specified with either '
 428                  '--sha1_file or --directory')
 429
 430   input_filename = args[0]
 431
 432   # Set output filename if not specified.
 433   if not options.output and not options.directory:
 434     if not options.sha1_file:
 435       # Target is a sha1 sum, so output filename would also be the sha1 sum.
 436       options.output = input_filename
 437     elif options.sha1_file:
 438       # Target is a .sha1 file.
 439       if not input_filename.endswith('.sha1'):
 440         parser.error('--sha1_file is specified, but the input filename '
 441                      'does not end with .sha1, and no --output is specified. '
 442                      'Either make sure the input filename has a .sha1 '
 443                      'extension, or specify --output.')
 444       options.output = input_filename[:-5]
 445     else:
 446       parser.error('Unreachable state.')
 447
 448   # Check if output file already exists.
 449   if not options.directory and not options.force and not options.no_resume:
 450     if os.path.exists(options.output):
 451       parser.error('Output file %s exists and --no_resume is specified.'
 452                    % options.output)
 453
 454   base_url = 'gs://%s' % options.bucket
 455
 456   # Check we have a valid bucket with valid permissions.
 457   if not options.no_auth:
 458     code = check_bucket_permissions(base_url, gsutil)
 459     if code:
 460       return code
 461
 462   return download_from_google_storage(
 463       input_filename, base_url, gsutil, options.num_threads, options.directory,
 464       options.recursive, options.force, options.output, options.ignore_errors,
 465       options.sha1_file, options.verbose, options.auto_platform)
 466
 467
 468 if __name__ == '__main__':
 469   sys.exit(main(sys.argv))