Reland 0c7d94eb9d9fa388bda0d74405c0928dceea22ff: Add support for tar.gz archive files to download f…

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 shutil
 import stat
 import sys
+import tarfile
 import threading
 import time
 
 import subprocess2
 
 
 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 = {
@@ skipping 16 matching lines @@
   pass
 
 
 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."""
   def __init__(self, path, boto_path=None, timeout=None, version='4.7'):
     if not os.path.exists(path):
       raise FileNotFoundError('GSUtil not found in %s' % path)
     self.path = path
     self.timeout = timeout
     self.boto_path = boto_path
@@ skipping 116 matching lines @@
           work_queue.put(
               (sha1_match.groups(1)[0], full_path.replace('.sha1', '')))
           work_queue_size += 1
         else:
           if not ignore_errors:
             raise InvalidFileError('No sha1 sum found in %s.' % filename)
           print >> sys.stderr, 'No sha1 sum found in %s.' % filename
   return work_queue_size
 
 
+def _validate_tar_file(tar, prefix):
+  def _validate(tarinfo):
+    """Returns false if the tarinfo is something we explicitly forbid."""
+    if tarinfo.issym() or tarinfo.islnk():
+      return False
+    if '..' in tarinfo.name or not tarinfo.name.startswith(prefix):
+      return False
+    return True
+  return all(map(_validate, tar.getmembers()))
+
 def _downloader_worker_thread(thread_num, q, force, base_url,
-                              gsutil, out_q, ret_codes, verbose):
+                              gsutil, out_q, ret_codes, verbose, extract,
+                              delete=True):
   while True:
     input_sha1_sum, output_filename = q.get()
     if input_sha1_sum is None:
       return
     if os.path.exists(output_filename) and not force:
       if get_sha1(output_filename) == input_sha1_sum:
         if verbose:
           out_q.put(
               '%d> File %s exists and SHA1 matches. Skipping.' % (
                   thread_num, output_filename))
@@ skipping 10 matching lines @@
       else:
         # Other error, probably auth related (bad ~/.boto, etc).
         out_q.put('%d> Failed to fetch file %s for %s, skipping. [Err: %s]' % (
             thread_num, file_url, output_filename, err))
         ret_codes.put((1, 'Failed to fetch file %s for %s. [Err: %s]' % (
             file_url, output_filename, err)))
       continue
     # Fetch the file.
     out_q.put('%d> Downloading %s...' % (thread_num, output_filename))
     try:
-      os.remove(output_filename)  # Delete the file if it exists already.
+      if delete:
+        os.remove(output_filename)  # Delete the file if it exists already.
     except OSError:
       if os.path.exists(output_filename):
         out_q.put('%d> Warning: deleting %s failed.' % (
             thread_num, output_filename))
     code, _, err = gsutil.check_call('cp', file_url, output_filename)
     if code != 0:
       out_q.put('%d> %s' % (thread_num, err))
       ret_codes.put((code, err))
       continue
 
     remote_sha1 = get_sha1(output_filename)
     if remote_sha1 != input_sha1_sum:
       msg = ('%d> ERROR remote sha1 (%s) does not match expected sha1 (%s).' %
              (thread_num, remote_sha1, input_sha1_sum))
       out_q.put(msg)
       ret_codes.put((20, msg))
       continue
 
+    if extract:
+      if (not tarfile.is_tarfile(output_filename)
+          or not output_filename.endswith('.tar.gz')):
+        out_q.put('%d> Error: %s is not a tar.gz archive.' % (
+                  thread_num, output_filename))
+        ret_codes.put((1, '%s is not a tar.gz archive.' % (output_filename)))
+        continue
+      with tarfile.open(output_filename, 'r:gz') as tar:
+        dirname = os.path.dirname(os.path.abspath(output_filename))
+        extract_dir = output_filename[0:len(output_filename)-7]
+        if not _validate_tar_file(tar, os.path.basename(extract_dir)):
+          out_q.put('%d> Error: %s contains files outside %s.' % (
+                    thread_num, output_filename, extract_dir))
+          ret_codes.put((1, '%s contains invalid entries.' % (output_filename)))
+          continue
+        if os.path.exists(extract_dir):
+          try:
+            shutil.rmtree(extract_dir)
+            out_q.put('%d> Removed %s...' % (thread_num, extract_dir))
+          except OSError:
+            out_q.put('%d> Warning: Can\'t delete: %s' % (
+                      thread_num, extract_dir))
+            ret_codes.put((1, 'Can\'t delete %s.' % (extract_dir)))
+            continue
+        out_q.put('%d> Extracting %d entries from %s to %s' %
+                  (thread_num, len(tar.getmembers()),output_filename,
+                   extract_dir))
+        tar.extractall(path=dirname)
     # Set executable bit.
     if sys.platform == 'cygwin':
       # Under cygwin, mark all files as executable. The executable flag in
       # Google Storage will not be set when uploading from Windows, so if
       # this script is running under cygwin and we're downloading an
       # executable, it will be unrunnable from inside cygwin without this.
       st = os.stat(output_filename)
       os.chmod(output_filename, st.st_mode | stat.S_IEXEC)
     elif sys.platform != 'win32':
       # On non-Windows platforms, key off of the custom header
@@ skipping 10 matching lines @@
   while True:
     line = output_queue.get()
     # Its plausible we want to print empty lines.
     if line is None:
       break
     print line
 
 
 def download_from_google_storage(
     input_filename, base_url, gsutil, num_threads, directory, recursive,
-    force, output, ignore_errors, sha1_file, verbose, auto_platform):
+    force, output, ignore_errors, sha1_file, verbose, auto_platform, extract):
   # Start up all the worker threads.
   all_threads = []
   download_start = time.time()
   stdout_queue = Queue.Queue()
   work_queue = Queue.Queue()
   ret_codes = Queue.Queue()
   ret_codes.put((0, None))
   for thread_num in range(num_threads):
     t = threading.Thread(
         target=_downloader_worker_thread,
         args=[thread_num, work_queue, force, base_url,
-              gsutil, stdout_queue, ret_codes, verbose])
+              gsutil, stdout_queue, ret_codes, verbose, extract])
     t.daemon = True
     t.start()
     all_threads.append(t)
   printer_thread = threading.Thread(target=printer_worker, args=[stdout_queue])
   printer_thread.daemon = True
   printer_thread.start()
 
   # Enumerate our work queue.
   work_queue_size = enumerate_work_queue(
       input_filename, work_queue, directory, recursive,
@@ skipping 67 matching lines @@
   parser.add_option('-p', '--platform',
                     help='A regular expression that is compared against '
                          'Python\'s sys.platform. If this option is specified, '
                          'the download will happen only if there is a match.')
   parser.add_option('-a', '--auto_platform',
                     action='store_true',
                     help='Detects if any parent folder of the target matches '
                          '(linux|mac|win).  If so, the script will only '
                          'process files that are in the paths that '
                          'that matches the current platform.')
+  parser.add_option('-u', '--extract',
+                    action='store_true',
+                    help='Extract a downloaded tar.gz file. '
+                         'Leaves the tar.gz file around for sha1 verification'
+                         'If a directory with the same name as the tar.gz '
+                         'file already exists, is deleted (to get a '
+                         'clean state in case of update.)')
   parser.add_option('-v', '--verbose', action='store_true',
                     help='Output extra diagnostic and progress information.')
 
   (options, args) = parser.parse_args()
 
   # Make sure we should run at all based on platform matching.
   if options.platform:
     if options.auto_platform:
       parser.error('--platform can not be specified with --auto_platform')
     if not re.match(options.platform, GetNormalizedPlatform()):
@@ skipping 73 matching lines @@
   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)
 
   base_url = 'gs://%s' % options.bucket
 
   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, options.verbose, options.auto_platform)
+      options.sha1_file, options.verbose, options.auto_platform,
+      options.extract)
 
 
 if __name__ == '__main__':
   sys.exit(main(sys.argv))