1 # Copyright (c) 2015, Google Inc. 2 # 3 # Permission to use, copy, modify, and/or distribute this software for any 4 # purpose with or without fee is hereby granted, provided that the above 5 # copyright notice and this permission notice appear in all copies. 6 # 7 # THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES 8 # WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF 9 # MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY 10 # SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES 11 # WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION 12 # OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN 13 # CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. 14 15 """Extracts archives.""" 16 17 18 import optparse 19 import os 20 import os.path 21 import tarfile 22 import shutil 23 import sys 24 import zipfile 25 26 27 def CheckedJoin(output, path): 28 """ 29 CheckedJoin returns os.path.join(output, path). It does sanity checks to 30 ensure the resulting path is under output, but shouldn't be used on untrusted 31 input. 32 """ 33 path = os.path.normpath(path) 34 if os.path.isabs(path) or path.startswith('.'): 35 raise ValueError(path) 36 return os.path.join(output, path) 37 38 39 def IterateZip(path): 40 """ 41 IterateZip opens the zip file at path and returns a generator of 42 (filename, mode, fileobj) tuples for each file in it. 43 """ 44 with zipfile.ZipFile(path, 'r') as zip_file: 45 for info in zip_file.infolist(): 46 if info.filename.endswith('/'): 47 continue 48 yield (info.filename, None, zip_file.open(info)) 49 50 51 def IterateTar(path): 52 """ 53 IterateTar opens the tar.gz file at path and returns a generator of 54 (filename, mode, fileobj) tuples for each file in it. 55 """ 56 with tarfile.open(path, 'r:gz') as tar_file: 57 for info in tar_file: 58 if info.isdir(): 59 continue 60 if not info.isfile(): 61 raise ValueError('Unknown entry type "%s"' % (info.name, )) 62 yield (info.name, info.mode, tar_file.extractfile(info)) 63 64 65 def main(args): 66 parser = optparse.OptionParser(usage='Usage: %prog ARCHIVE OUTPUT') 67 parser.add_option('--no-prefix', dest='no_prefix', action='store_true', 68 help='Do not remove a prefix from paths in the archive.') 69 options, args = parser.parse_args(args) 70 71 if len(args) != 2: 72 parser.print_help() 73 return 1 74 75 archive, output = args 76 77 if not os.path.exists(archive): 78 # Skip archives that weren't downloaded. 79 return 0 80 81 if archive.endswith('.zip'): 82 entries = IterateZip(archive) 83 elif archive.endswith('.tar.gz'): 84 entries = IterateTar(archive) 85 else: 86 raise ValueError(archive) 87 88 try: 89 if os.path.exists(output): 90 print "Removing %s" % (output, ) 91 shutil.rmtree(output) 92 93 print "Extracting %s to %s" % (archive, output) 94 prefix = None 95 num_extracted = 0 96 for path, mode, inp in entries: 97 # Even on Windows, zip files must always use forward slashes. 98 if '\\' in path or path.startswith('/'): 99 raise ValueError(path) 100 101 if not options.no_prefix: 102 new_prefix, rest = path.split('/', 1) 103 104 # Ensure the archive is consistent. 105 if prefix is None: 106 prefix = new_prefix 107 if prefix != new_prefix: 108 raise ValueError((prefix, new_prefix)) 109 else: 110 rest = path 111 112 # Extract the file into the output directory. 113 fixed_path = CheckedJoin(output, rest) 114 if not os.path.isdir(os.path.dirname(fixed_path)): 115 os.makedirs(os.path.dirname(fixed_path)) 116 with open(fixed_path, 'wb') as out: 117 shutil.copyfileobj(inp, out) 118 119 # Fix up permissions if needbe. 120 # TODO(davidben): To be extra tidy, this should only track the execute bit 121 # as in git. 122 if mode is not None: 123 os.chmod(fixed_path, mode) 124 125 # Print every 100 files, so bots do not time out on large archives. 126 num_extracted += 1 127 if num_extracted % 100 == 0: 128 print "Extracted %d files..." % (num_extracted,) 129 finally: 130 entries.close() 131 132 if num_extracted % 100 == 0: 133 print "Done. Extracted %d files." % (num_extracted,) 134 135 return 0 136 137 138 if __name__ == '__main__': 139 sys.exit(main(sys.argv[1:])) 140