Scripts to download files from google storage based on sha1 sums

build/download_from_google_storage.py

Index: build/download_from_google_storage.py
diff --git a/build/download_from_google_storage.py b/build/download_from_google_storage.py
new file mode 100755
index 0000000000000000000000000000000000000000..6a5db8b1f7cc3c210675a0e56137e6e2a1e16dcf
--- /dev/null
+++ b/build/download_from_google_storage.py
@@ -0,0 +1,239 @@
+#!/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.
+
+"""Script to download files from Google Storage."""
+
+
+import hashlib
+import os
+import Queue
+import re
+import subprocess
+import sys
+import tempfile
+import threading
+import time
+import zipfile
+
+from optparse import OptionParser

[M-A Ruel, 2013/01/10 02:33:08]

Actually, why import this one differently from the others? Just "import
optparse" like the other stdlibs.

[Ryan Tseng, 2013/01/14 21:37:13]

Done.

+
+GSUTIL_DEFAULT_PATH = os.path.join(os.path.dirname(os.path.normpath(__file__)),
+          '..', '..', 'third_party', 'gsutil', 'gsutil')

[M-A Ruel, 2013/01/10 02:33:08]

Aligning at +10 seems quite arbitrary.

[Ryan Tseng, 2013/01/14 21:37:13]

Fixed

+
+
+class Gsutil():
+  def __init__(self, path, boto_path=None, timeout=None):
+    if os.path.exists(path):

[M-A Ruel, 2013/01/10 02:33:08]

if not os.path.exists(path):
  raise OSError('GSUtil not found in %s' % path)
self.path = path

[Ryan Tseng, 2013/01/14 21:37:13]

Done.

+      self.path = path
+    else:
+      raise OSError('GSUtil not found in %s' % path)
+    self.timeout = timeout
+    self.boto_path = boto_path
+
+  def call(self, *args):
+    def _thread_main():
+      thr = threading.current_thread()
+      env = os.environ.copy()
+      if self.boto_path is not None:
+        env['AWS_CREDENTIAL_FILE'] = self.boto_path
+      p = subprocess.Popen((sys.executable, self.path) + args, env=env)

[M-A Ruel, 2013/01/10 02:33:08]

you want subprocess.call()

[Ryan Tseng, 2013/01/14 21:37:13]

Done.

+      thr.status = p.wait()
+    t = threading.Thread(target=_thread_main)
+    t.start()
+    t.join(self.timeout)
+    if thr.isAlive():
+      raise RuntimeError('%s %s timed out after %d seconds.' % (
+              self.path, ' '.join(args), self.timeout))
+    return thr.status
+
+  def check_call(self, *args):
+    def _thread_main():

[M-A Ruel, 2013/01/10 02:33:08]

This code is never called?

[Ryan Tseng, 2013/01/14 21:37:13]

Done.

+      thr = threading.current_thread()
+      env = os.environ.copy()
+      if self.boto_path is not None:
+        env['AWS_CREDENTIAL_FILE'] = self.boto_path
+      p = subprocess.Popen((sys.executable, self.path) + args,
+                           stdout=subprocess.PIPE,
+                           stderr=subprocess.PIPE,
+                           env=env)
+    thr.status = p.wait()
+    out, err = p.communicate()
+
+    if code == 0:
+      return 0
+
+    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 CheckSHA1(sha1_sum, filename):
+  sha1 = hashlib.sha1()

[M-A Ruel, 2013/01/10 02:33:08]

FYI, this doesn't work with files > 1.5 gb or so. See the implementation in
swarm_client/ for one that works with larger files.

[Ryan Tseng, 2013/01/14 21:37:13]

Done.

+  sha1.update(open(filename).read())
+  return sha1_sum == sha1.hexdigest()
+
+
+def _downloader_worker_thread(thread_num, q, options, base_url, gsutil):
+  while True:
+    try:
+      input_sha1_sum, output_filename = q.get_nowait()
+      if os.path.exists(output_filename) and not options.force:
+        if CheckSHA1(input_sha1_sum, output_filename):
+          print '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:
+        print >>sys.stderr, 'File %s for %s does not exist, skipping.' % (
+            file_url, output_filename)
+        continue
+      # Fetch the file.
+      print 'Downloading %s to %s...' % (file_url, output_filename)
+      code, out, err = gsutil.call('cp', '-q', file_url, output_filename)
+      if code != 0:
+        print >>sys.stderr, gsutil.stderr
+        return code
+    except Queue.Empty:
+      return
+
+
+def main(args):
+  usage = ('usage: %prog [options] target\nTarget must be:\n'
+           '(1) a directory.\n(2) a sha1 sum ([A-Za-z0-9]{40}).\n'
+           '(3) a .sha1 file, containing a sha1 sum on the first line.')
+  parser = OptionParser(usage)
+  parser.add_option('-o', '--output', default=None,
+                    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', default='chrome-artifacts',
+                    help='Google Storage bucket to fetch from.')
+  parser.add_option('-f', '--force', action='store_true', default=False,
+                    help='Force download even if local file exists.')
+  parser.add_option('-r', '--recursive', action='store_true', default=False,
+                    help='Scan folders recursively for .sha1 files.')
+  parser.add_option('-t', '--num_threads', default=1, type='int',
+                    help='Number of downloader threads to run.')
+  # This file should be stored in tools/deps_scripts/ and we want the path to
+  # third_party/gsutil/gsutil
+  parser.add_option('-g', '--gsutil_path', default=GSUTIL_DEFAULT_PATH,
+          help='Path to the gsutil script.')
+
+  (options, args) = parser.parse_args()
+  if len(args) < 1:
+    print >>sys.stderr, 'ERROR: Missing target.'

[M-A Ruel, 2013/01/10 02:33:08]

parser.error() and line 142 too

[Ryan Tseng, 2013/01/14 21:37:13]

Done.

+    parser.print_help()
+    return 1
+  elif len(args) > 1:
+    # TODO(hinoka): Multi target support.
+    print >>sys.stderr, 'ERROR: Too many targets.'
+    parser.print_help()
+    return 1
+  else:
+    input_filename = args[0]
+
+  # input_filename is a file? This could mean one of three things:

[M-A Ruel, 2013/01/10 02:33:08]

Remove the guess work and require an argument.

Guesswork = maintenance nightmare in a few years.

[Ryan Tseng, 2013/01/14 21:37:13]

Done.

+  #   1. The input is a directory
+  #   2. The input is a .sha1 file
+  #   3. The input is an already downloaded binary file.
+  if os.path.exists(input_filename):
+    if os.path.isdir(input_filename):
+      # Check if the input is a directory.
+      dir_name = True
+      checked_sha1 = False
+    else:
+      with open(input_filename) as f:
+        sha1_match = re.search('^([A-Za-z0-9]{40})\s*$', f.read(1024))
+      if sha1_match:
+        # Check if we can match a sha1 sum in the first 1024 bytes.
+        if input_filename.endswith('.sha1'):
+          options.output = input_filename[:-5]
+        input_filename = sha1_match.groups(1)
+        dir_name = False
+        checked_sha1 = False
+      elif CheckSHA1(input_filename, input_filename):
+        # Check if input_filename is already downloaded.
+        dir_name = False
+        checked_sha1 = True
+      else:
+        if not re.match('[A-Za-z0-9]{40}', input_filename):
+          print >>sys.stderr, 'Input %s not recognized.' % input_filename
+          parser.print_help()
+          return 1
+
+  if not options.output:
+    options.output = input_filename
+  base_url = 'gs://%s' % options.bucket
+
+  if os.path.exists(options.gsutil_path):
+    gsutil = Gsutil(options.gsutil_path)
+  else:
+    for path in os.environ["PATH"].split(os.pathsep):
+      if os.path.exists(path) and 'gsutil' in os.listdir(path):
+        gsutil = Gsutil(os.path.join(path, 'gsutil'))
+
+  # Check if we have permissions.
+  code, ls_out, 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, exiting' % base_url
+      return 403
+  elif code == 404:
+    print >>sys.stderr, '%s not found.' % base_url
+    return 404
+  elif code != 0:
+    print >>sys.stderr, ls_err
+    return code
+
+  # Enumerate our work queue.
+  work_queue = Queue.Queue()
+  if dir_name:
+    if options.recursive:
+      for root, dirs, files in os.walk(input_filename):
+        if '.svn' in dirs:
+          dirs.remove('.svn')
+        if not options.recursive:
+          for item in dirs:
+            dirs.remove(item)
+        for filename in files:
+          full_path = os.path.join(root, filename)
+          if full_path.endswith('.sha1'):
+            with open(full_path) as f:
+              sha1_match = re.search('([A-Za-z0-9]{40})', f.read(1024))
+            if sha1_match:
+              work_queue.put((sha1_match.groups(1)[0],
+                              full_path.replace('.sha1', '')))
+  else:
+    work_queue.put((input_filename, options.output))
+
+  # Start up all the worker threads.
+  all_threads = []
+  for thread_num in range(options.num_threads):
+    t = threading.Thread(target=_downloader_worker_thread, args=[thread_num,

[M-A Ruel, 2013/01/10 02:33:08]

Are you going to start 1000 threads if there are 1000 files to download?

[Ryan Tseng, 2013/01/14 21:37:13]

Nope, it'll start 1 thread by default, or 10 threads if "-t 10" is specified,
and do 10 at a time.

+        work_queue, options, base_url, gsutil.clone()])
+    t.daemon = True
+    t.start()
+    all_threads.append(t)
+
+  # Wait for all downloads to finish.
+  for t in threads:
+    t.join()

[M-A Ruel, 2013/01/10 02:33:08]

return 0

+
+
+if __name__ == '__main__':
+  sys.exit(main(sys.argv))