Scripts to download files from google storage based on sha1 sums

download_from_google_storage.py

+#!/usr/bin/env python
+# Copyright (c) 2012 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 based on SHA1 sums."""
+
+
+import hashlib
+import optparse
+import os
+import Queue
+import re
+import sys
+import threading
+import time
+
+import subprocess2
+
+
+GSUTIL_DEFAULT_PATH = os.path.join(
+    os.path.dirname(os.path.abspath(__file__)),
+    'third_party', 'gsutil', 'gsutil')
+
+
+# Common utilities
+class Gsutil(object):
+  """Call gsutil with some predefined settings."""
+  def __init__(self, path, boto_path=None, timeout=None):
+    if not os.path.exists(path):
+      raise OSError('GSUtil not found in %s' % path)
+    self.path = path
+    self.timeout = timeout
+    self.boto_path = boto_path
+
+  def call(self, *args):
+    env = os.environ.copy()
+    if self.boto_path is not None:
+      env['AWS_CREDENTIAL_FILE'] = self.boto_path
+    return subprocess2.call((sys.executable, self.path) + args,
+                            env=env,
+                            timeout=self.timeout)
+
+  def check_call(self, *args):
+    env = os.environ.copy()
+    if self.boto_path is not None:
+      env['AWS_CREDENTIAL_FILE'] = self.boto_path
+    ((out, err), code) = subprocess2.communicate(
+        (sys.executable, self.path) + args,
+        stdout=subprocess2.PIPE,
+        stderr=subprocess2.PIPE,
+        env=env,
+        timeout=self.timeout)
+
+    # Parse output.
+    status_code_match = re.search('status=([0-9]+)', err)
+    if status_code_match:
+      return int(status_code_match.groups(1))
+    elif ('You are attempting to access protected data with '
+          'no configured credentials.' in err):
+      return (403, out, err)
+    elif 'No such object' in err:
+      return (404, out, err)
+    else:
+      return (code, out, err)
+
+  def clone(self):
+    return Gsutil(self.path, self.boto_path, self.timeout)
+
+
+def CheckBucketPermissions(bucket, gsutil):
+  if not bucket:
+    print >> sys.stderr, 'Missing bucket %s.'
+    return (None, 1)
+  base_url = 'gs://%s' % bucket
+
+  # Check if we have permissions to the Google Storage bucket.

[M-A Ruel, 2013/03/06 19:49:59]

You can remove this comment.

[Ryan Tseng, 2013/03/06 20:19:08]

Done.

+  code, _, ls_err = gsutil.check_call('ls', base_url)
+  if code == 403:
+    code, _, _ = gsutil.call('config')
+    if code != 0:
+      print >> sys.stderr, 'Error while authenticating to %s.' % base_url
+  elif code == 404:
+    print >> sys.stderr, '%s not found.' % base_url
+  elif code != 0:
+    print >> sys.stderr, ls_err
+  return (base_url, code)
+
+
+def GetSHA1(filename):
+  sha1 = hashlib.sha1()
+  with open(filename, 'rb') as f:
+    while True:
+      # Read in 1mb chunks, so it doesn't all have to be loaded into memory.
+      chunk = f.read(1024*1024)
+      if not chunk:
+        break
+      sha1.update(chunk)
+  return sha1.hexdigest()
+
+
+# Download-specific code starts here
+
+def enumerate_work_queue(input_filename, work_queue, directory,
+                         recursive, ignore_errors, output, sha1_file):
+  if sha1_file:
+    if not os.path.exists(input_filename):
+      print >> sys.stderr, '%s not found.' % input_filename
+      if not ignore_errors:
+        raise Exception('%s not found.' % input_filename)
+    with open(input_filename, 'rb') as f:
+      sha1_match = re.match('^([A-Za-z0-9]{40})$', f.read(1024).rstrip())
+      if sha1_match:
+        work_queue.put(
+            (sha1_match.groups(1)[0], input_filename.replace('.sha1', '')))
+        return 1
+    print >> sys.stderr, 'No sha1 sum found in %s.' % input_filename
+    if not ignore_errors:
+      raise Exception('No sha1 sum found in %s.' % input_filename)
+    return 0
+
+  if not directory:
+    work_queue.put((input_filename, output))
+    return 1
+
+  work_queue_size = 0
+  for root, dirs, files in os.walk(input_filename):
+    if not recursive:
+      for item in dirs[:]:
+        dirs.remove(item)
+    else:
+      for exclude in ['.svn', '.git']:
+        if exclude in dirs:
+          dirs.remove(exclude)
+    for filename in files:
+      full_path = os.path.join(root, filename)
+      if full_path.endswith('.sha1'):
+        with open(full_path, 'rb') as f:
+          sha1_match = re.match('^([A-Za-z0-9]{40})$', f.read(1024).rstrip())
+        if sha1_match:
+          work_queue.put(
+              (sha1_match.groups(1)[0], full_path.replace('.sha1', '')))
+          work_queue_size += 1
+        else:
+          print >> sys.stderr, 'No sha1 sum found in %s.' % filename
+          if not ignore_errors:
+            raise Exception('No sha1 sum found in %s.' % filename)
+  return work_queue_size
+
+
+def _downloader_worker_thread(thread_num, q, force, base_url, gsutil, out_q):
+  while True:
+    input_sha1_sum, output_filename = q.get()
+    if input_sha1_sum is None:
+      out_q.put('Thread %d is done' % thread_num)
+      return
+    if os.path.exists(output_filename) and not force:
+      if GetSHA1(output_filename) == input_sha1_sum:
+        out_q.put(
+            'File %s exists and SHA1 sum (%s) matches. Skipping.' % (
+                output_filename , input_sha1_sum))
+        continue
+    # Check if file exists.
+    file_url = '%s/%s' % (base_url, input_sha1_sum)
+    if gsutil.check_call('ls', file_url)[0] != 0:
+      out_q.put('File %s for %s does not exist, skipping.' % (
+          file_url, output_filename))
+      continue
+    # Fetch the file.
+    out_q.put('Downloading %s to %s...' % (file_url, output_filename))
+    code, _, err = gsutil.check_call('cp', '-q', file_url, output_filename)
+    if code != 0:
+      out_q.put(err)
+      return code
+
+
+def download_from_google_storage(
+    input_filename, base_url, gsutil, num_threads, directory, recursive,
+    force, output, ignore_errors, sha1_file):
+  # Start up all the worker threads.
+  all_threads = []
+  download_timer = time.time()
+  stdout_queue = Queue.Queue()
+  work_queue = Queue.Queue()
+  for thread_num in range(num_threads):
+    t = threading.Thread(
+        target=_downloader_worker_thread,
+        args=[thread_num, work_queue, force, base_url,
+              gsutil.clone(), stdout_queue])
+    t.daemon = True
+    t.start()
+    all_threads.append(t)
+
+  # Enumerate our work queue.
+  work_queue_size = enumerate_work_queue(
+      input_filename, work_queue, directory, recursive,
+      ignore_errors, output, sha1_file)
+  for _ in all_threads:
+    work_queue.put((None, None))  # Used to tell worker threads to stop.
+
+  # Wait for all downloads to finish.
+  while not work_queue.empty() and any(t.is_alive() for t in all_threads):

[M-A Ruel, 2013/03/06 19:49:59]

There's a possibility that the code would get there and no entry would have been
appended to work_queue yet. It's possible on a single-cpu system with a single
large file. Then this check would immediately fall through line 204 and could
possibly miss items. This need to be reworked to be race-condition free.

[Ryan Tseng, 2013/03/06 20:19:08]

isn't appending to a Queue() synchronous?  If we pass lines 195 and 199 can't we
guarantee that the work queue isn't empty?

[Ryan Tseng, 2013/03/06 20:24:45]

oh nevermind I see what you're saying.

+    print stdout_queue.get()
+  while not stdout_queue.empty():
+    print stdout_queue.get()
+
+  print 'Success.'
+  print 'Downloading %d files took %1f second(s)' % (
+      work_queue_size, time.time() - download_timer)
+  return 0
+
+
+def main(args):
+  usage = ('usage: %prog [options] target\nTarget must be:\n'
+           '(default) a sha1 sum ([A-Za-z0-9]{40}).\n(-s or --sha1_file) a '
+           '.sha1 file, containing a sha1 sum on the first line. (-d or '
+           '--directory) A directory to scan for .sha1 files. ')
+  parser = optparse.OptionParser(usage)
+  parser.add_option('-o', '--output',
+                    help='Specify the output file name. Defaults to:\n'
+                         '(a) Given a SHA1 hash, the name is the SHA1 hash.\n'
+                         '(b) Given a .sha1 file or directory, the name will '
+                         'match (.*).sha1.')
+  parser.add_option('-b', '--bucket',
+                    help='Google Storage bucket to fetch from.')
+  parser.add_option('-e', '--boto',
+                    help='Specify a custom boto file.')
+  parser.add_option('-c', '--no_resume', action='store_true',
+                    help='Resume download if file is partially downloaded.')
+  parser.add_option('-f', '--force', action='store_true',
+                    help='Force download even if local file exists.')
+  parser.add_option('-i', '--ignore_errors', action='store_true',
+                    help='Don\'t throw error if we find an invalid .sha1 file.')
+  parser.add_option('-r', '--recursive', action='store_true',
+                    help='Scan folders recursively for .sha1 files. '
+                         'Must be used with -d/--directory')
+  parser.add_option('-t', '--num_threads', default=1, type='int',
+                    help='Number of downloader threads to run.')
+  parser.add_option('-d', '--directory', action='store_true',
+                    help='The target is a directory.  '
+                         'Cannot be used with -s/--sha1_file.')
+  parser.add_option('-s', '--sha1_file', action='store_true',
+                    help='The target is a file containing a sha1 sum.  '
+                         'Cannot be used with -d/--directory.')
+
+  (options, args) = parser.parse_args()
+  if not args:
+    parser.error('Missing target.')
+  if len(args) > 1:
+    parser.error('Too many targets.')
+  if not options.bucket:
+    parser.error('Missing bucket.  Specify bucket with --bucket.')
+  if options.sha1_file and options.directory:
+    parser.error('Both --directory and --sha1_file are specified, '
+                 'can only specify one.')
+  elif options.recursive and not options.directory:
+    parser.error('--recursive specified but --directory not specified.')
+  elif options.output and options.directory:
+    parser.error('--directory is specified, so --output has no effect.')
+  else:
+    input_filename = args[0]
+
+  # Set output filename if not specified.
+  if not options.output and not options.directory:
+    if not options.sha1_file:
+      # Target is a sha1 sum, so output filename would also be the sha1 sum.
+      options.output = input_filename
+    elif options.sha1_file:
+      # Target is a .sha1 file.
+      if not input_filename.endswith('.sha1'):
+        parser.error('--sha1_file is specified, but the input filename '
+                     'does not end with .sha1, and no --output is specified. '
+                     'Either make sure the input filename has a .sha1 '
+                     'extension, or specify --output.')
+      options.output = input_filename[:-5]
+    else:
+      raise NotImplementedError('Unreachable state.')
+
+  # Check if output file already exists.
+  if not options.directory and not options.force and not options.no_resume:
+    if os.path.exists(options.output):
+      parser.error('Output file %s exists and --no_resume is specified.'
+                   % options.output)
+
+  # Make sure we can find a working instance of gsutil.
+  if os.path.exists(GSUTIL_DEFAULT_PATH):
+    gsutil = Gsutil(GSUTIL_DEFAULT_PATH)
+  else:
+    print >> sys.stderr, ('gsutil not found in %s, bad depot_tools checkout?' %
+                          GSUTIL_DEFAULT_PATH)
+    return 1
+
+  # Check we have a valid bucket with valid permissions.
+  base_url, code = CheckBucketPermissions(options.bucket, gsutil)
+  if code:
+    return code
+
+  return download_from_google_storage(
+      input_filename, base_url, gsutil, options.num_threads, options.directory,
+      options.recursive, options.force, options.output, options.ignore_errors,
+      options.sha1_file)
+
+
+if __name__ == '__main__':
+  sys.exit(main(sys.argv))