Add support for tar.gz archive files to download from download_from_google_storage

download_from_google_storage.py

Index: download_from_google_storage.py
===================================================================
--- download_from_google_storage.py	(revision 293975)
+++ download_from_google_storage.py	(working copy)
@@ -11,8 +11,10 @@
 import os
 import Queue
 import re
+import shutil
 import stat
 import sys
+import tarfile
 import threading
 import time
 
@@ -49,7 +51,6 @@
     return 'win32'
   return sys.platform
 
-
 # Common utilities
 class Gsutil(object):
   """Call gsutil with some predefined settings.  This is a convenience object,
@@ -203,8 +204,12 @@
   return work_queue_size
 
 
+def _validate_tar_file(tar, prefix):
+  files = tar.getnames()
+  return reduce(lambda x, y: x + 1 if not y.startswith(prefix) else x, files, 0)

[hinoka, 2015/02/06 18:54:35]

1. Still prefer any(map(...))
2. This doesn't look like it'd verify something like
foo/bar/baz/../../../../../../../bin/bash.  You need to do some verification on
the os.realpath() version of the path.

[ricow1, 2015/02/10 11:15:45]

Added more checks

+
 def _downloader_worker_thread(thread_num, q, force, base_url,
-                              gsutil, out_q, ret_codes, verbose):
+                              gsutil, out_q, ret_codes, verbose, extract):
   while True:
     input_sha1_sum, output_filename = q.get()
     if input_sha1_sum is None:
@@ -237,6 +242,34 @@
       out_q.put('%d> %s' % (thread_num, err))
       ret_codes.put((code, err))
 
+    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
+      tar = tarfile.open(output_filename, 'r:gz')
+      dirname = os.path.dirname(os.path.abspath(output_filename))
+      extract_dir = output_filename[0:len(output_filename)-7]
+      if _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
+      out_q.put('%d> extracting %s...' % (thread_num, extract_dir))

[hinoka, 2015/02/06 18:54:35]

nit: s/extracting/Extracting/

[ricow1, 2015/02/10 11:15:45]

Done.

+      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 %s to %s' % (thread_num, 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
@@ -267,7 +300,7 @@
 
 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()
@@ -279,7 +312,7 @@
     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)
@@ -367,6 +400,13 @@
                          '(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 after download. '

[hinoka, 2015/02/06 18:54:35]

s/after download// (its redundant)

[ricow1, 2015/02/10 11:15:45]

Done.

+                         'Leaves the tar.gz file around for sha verification1'

[hinoka, 2015/02/06 18:54:35]

s/sha verification1/sha1 verification/

[ricow1, 2015/02/10 11:15:45]

Done.

+                         'If a directory with the same name as the tar.gz '
+                         'file already exists, this is deleted (to get a '

[hinoka, 2015/02/06 18:54:35]

s/this/it/

[ricow1, 2015/02/10 11:15:45]

Done.

+                         'clean state in case of update.)')
   parser.add_option('-v', '--verbose', action='store_true',
                     help='Output extra diagnostic and progress information.')
 
@@ -463,7 +503,8 @@
   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__':