blob: 2aa4bffc4281c6f305d581e82ed6bd9a2428e2b9 [file] [log] [blame]
Ronald Oussoren143cefb2006-06-15 08:14:18 +00001"""
2Read and write ZIP files.
3"""
Georg Brandl62416bc2008-01-07 18:47:44 +00004import struct, os, time, sys, shutil
Martin v. Löwis0dfcfc82009-01-24 14:00:33 +00005import binascii, cStringIO, stat
Guido van Rossum32abe6f2000-03-31 17:30:02 +00006
7try:
Tim Peterse1190062001-01-15 03:34:38 +00008 import zlib # We may need its compression method
Gregory P. Smithb89a0962008-03-19 01:46:10 +00009 crc32 = zlib.crc32
Guido van Rossum9c673f32001-04-10 15:37:12 +000010except ImportError:
Guido van Rossum32abe6f2000-03-31 17:30:02 +000011 zlib = None
Gregory P. Smithb89a0962008-03-19 01:46:10 +000012 crc32 = binascii.crc32
Guido van Rossum32abe6f2000-03-31 17:30:02 +000013
Skip Montanaro40fc1602001-03-01 04:27:19 +000014__all__ = ["BadZipfile", "error", "ZIP_STORED", "ZIP_DEFLATED", "is_zipfile",
Ronald Oussoren143cefb2006-06-15 08:14:18 +000015 "ZipInfo", "ZipFile", "PyZipFile", "LargeZipFile" ]
Skip Montanaro40fc1602001-03-01 04:27:19 +000016
Fred Drake5db246d2000-09-29 20:44:48 +000017class BadZipfile(Exception):
Guido van Rossum32abe6f2000-03-31 17:30:02 +000018 pass
Ronald Oussoren143cefb2006-06-15 08:14:18 +000019
20
21class LargeZipFile(Exception):
Tim Petersa608bb22006-06-15 18:06:29 +000022 """
Ronald Oussoren143cefb2006-06-15 08:14:18 +000023 Raised when writing a zipfile, the zipfile requires ZIP64 extensions
24 and those extensions are disabled.
25 """
26
Tim Peterse1190062001-01-15 03:34:38 +000027error = BadZipfile # The exception raised by this module
Guido van Rossum32abe6f2000-03-31 17:30:02 +000028
Amaury Forgeot d'Arcd25f87a2009-01-17 16:40:17 +000029ZIP64_LIMIT = (1 << 31) - 1
Martin v. Löwis8c436412008-07-03 12:51:14 +000030ZIP_FILECOUNT_LIMIT = 1 << 16
31ZIP_MAX_COMMENT = (1 << 16) - 1
Ronald Oussoren143cefb2006-06-15 08:14:18 +000032
Guido van Rossum32abe6f2000-03-31 17:30:02 +000033# constants for Zip file compression methods
34ZIP_STORED = 0
35ZIP_DEFLATED = 8
36# Other ZIP compression methods not supported
37
Martin v. Löwis8c436412008-07-03 12:51:14 +000038# Below are some formats and associated data for reading/writing headers using
39# the struct module. The names and structures of headers/records are those used
40# in the PKWARE description of the ZIP file format:
41# http://www.pkware.com/documents/casestudies/APPNOTE.TXT
42# (URL valid as of January 2008)
Ronald Oussoren143cefb2006-06-15 08:14:18 +000043
Martin v. Löwis8c436412008-07-03 12:51:14 +000044# The "end of central directory" structure, magic number, size, and indices
45# (section V.I in the format document)
Amaury Forgeot d'Arcae6d2b92008-07-11 21:28:25 +000046structEndArchive = "<4s4H2LH"
47stringEndArchive = "PK\005\006"
48sizeEndCentDir = struct.calcsize(structEndArchive)
Martin v. Löwis8c436412008-07-03 12:51:14 +000049
50_ECD_SIGNATURE = 0
51_ECD_DISK_NUMBER = 1
52_ECD_DISK_START = 2
53_ECD_ENTRIES_THIS_DISK = 3
54_ECD_ENTRIES_TOTAL = 4
55_ECD_SIZE = 5
56_ECD_OFFSET = 6
57_ECD_COMMENT_SIZE = 7
58# These last two indices are not part of the structure as defined in the
59# spec, but they are used internally by this module as a convenience
60_ECD_COMMENT = 8
61_ECD_LOCATION = 9
62
63# The "central directory" structure, magic number, size, and indices
64# of entries in the structure (section V.F in the format document)
65structCentralDir = "<4s4B4HL2L5H2L"
Amaury Forgeot d'Arcae6d2b92008-07-11 21:28:25 +000066stringCentralDir = "PK\001\002"
Martin v. Löwis8c436412008-07-03 12:51:14 +000067sizeCentralDir = struct.calcsize(structCentralDir)
68
Fred Drake3e038e52001-02-28 17:56:26 +000069# indexes of entries in the central directory structure
70_CD_SIGNATURE = 0
71_CD_CREATE_VERSION = 1
72_CD_CREATE_SYSTEM = 2
73_CD_EXTRACT_VERSION = 3
Martin v. Löwis8c436412008-07-03 12:51:14 +000074_CD_EXTRACT_SYSTEM = 4
Fred Drake3e038e52001-02-28 17:56:26 +000075_CD_FLAG_BITS = 5
76_CD_COMPRESS_TYPE = 6
77_CD_TIME = 7
78_CD_DATE = 8
79_CD_CRC = 9
80_CD_COMPRESSED_SIZE = 10
81_CD_UNCOMPRESSED_SIZE = 11
82_CD_FILENAME_LENGTH = 12
83_CD_EXTRA_FIELD_LENGTH = 13
84_CD_COMMENT_LENGTH = 14
85_CD_DISK_NUMBER_START = 15
86_CD_INTERNAL_FILE_ATTRIBUTES = 16
87_CD_EXTERNAL_FILE_ATTRIBUTES = 17
88_CD_LOCAL_HEADER_OFFSET = 18
89
Martin v. Löwis8c436412008-07-03 12:51:14 +000090# The "local file header" structure, magic number, size, and indices
91# (section V.A in the format document)
92structFileHeader = "<4s2B4HL2L2H"
Amaury Forgeot d'Arcae6d2b92008-07-11 21:28:25 +000093stringFileHeader = "PK\003\004"
Martin v. Löwis8c436412008-07-03 12:51:14 +000094sizeFileHeader = struct.calcsize(structFileHeader)
95
Fred Drake3e038e52001-02-28 17:56:26 +000096_FH_SIGNATURE = 0
97_FH_EXTRACT_VERSION = 1
Martin v. Löwis8c436412008-07-03 12:51:14 +000098_FH_EXTRACT_SYSTEM = 2
Fred Drake3e038e52001-02-28 17:56:26 +000099_FH_GENERAL_PURPOSE_FLAG_BITS = 3
100_FH_COMPRESSION_METHOD = 4
101_FH_LAST_MOD_TIME = 5
102_FH_LAST_MOD_DATE = 6
103_FH_CRC = 7
104_FH_COMPRESSED_SIZE = 8
105_FH_UNCOMPRESSED_SIZE = 9
106_FH_FILENAME_LENGTH = 10
107_FH_EXTRA_FIELD_LENGTH = 11
108
Martin v. Löwis8c436412008-07-03 12:51:14 +0000109# The "Zip64 end of central directory locator" structure, magic number, and size
Amaury Forgeot d'Arcae6d2b92008-07-11 21:28:25 +0000110structEndArchive64Locator = "<4sLQL"
111stringEndArchive64Locator = "PK\x06\x07"
112sizeEndCentDir64Locator = struct.calcsize(structEndArchive64Locator)
Martin v. Löwis8c436412008-07-03 12:51:14 +0000113
114# The "Zip64 end of central directory" record, magic number, size, and indices
115# (section V.G in the format document)
Amaury Forgeot d'Arcae6d2b92008-07-11 21:28:25 +0000116structEndArchive64 = "<4sQ2H2L4Q"
117stringEndArchive64 = "PK\x06\x06"
118sizeEndCentDir64 = struct.calcsize(structEndArchive64)
Martin v. Löwis8c436412008-07-03 12:51:14 +0000119
120_CD64_SIGNATURE = 0
121_CD64_DIRECTORY_RECSIZE = 1
122_CD64_CREATE_VERSION = 2
123_CD64_EXTRACT_VERSION = 3
124_CD64_DISK_NUMBER = 4
125_CD64_DISK_NUMBER_START = 5
126_CD64_NUMBER_ENTRIES_THIS_DISK = 6
127_CD64_NUMBER_ENTRIES_TOTAL = 7
128_CD64_DIRECTORY_SIZE = 8
129_CD64_OFFSET_START_CENTDIR = 9
130
Antoine Pitrou6f193e02008-12-27 15:43:12 +0000131def _check_zipfile(fp):
Guido van Rossum32abe6f2000-03-31 17:30:02 +0000132 try:
Antoine Pitrou6f193e02008-12-27 15:43:12 +0000133 if _EndRecData(fp):
134 return True # file has correct magic number
Fred Drake7e473802001-05-11 19:52:57 +0000135 except IOError:
Guido van Rossum32abe6f2000-03-31 17:30:02 +0000136 pass
Guido van Rossum8ca162f2002-04-07 06:36:23 +0000137 return False
Guido van Rossum32abe6f2000-03-31 17:30:02 +0000138
Antoine Pitrou6f193e02008-12-27 15:43:12 +0000139def is_zipfile(filename):
140 """Quickly see if a file is a ZIP file by checking the magic number.
141
142 The filename argument may be a file or file-like object too.
143 """
144 result = False
145 try:
146 if hasattr(filename, "read"):
147 result = _check_zipfile(fp=filename)
148 else:
149 with open(filename, "rb") as fp:
150 result = _check_zipfile(fp)
151 except IOError:
152 pass
153 return result
154
Ronald Oussoren143cefb2006-06-15 08:14:18 +0000155def _EndRecData64(fpin, offset, endrec):
156 """
157 Read the ZIP64 end-of-archive records and use that to update endrec
158 """
Martin v. Löwis8c436412008-07-03 12:51:14 +0000159 fpin.seek(offset - sizeEndCentDir64Locator, 2)
160 data = fpin.read(sizeEndCentDir64Locator)
Amaury Forgeot d'Arcae6d2b92008-07-11 21:28:25 +0000161 sig, diskno, reloff, disks = struct.unpack(structEndArchive64Locator, data)
162 if sig != stringEndArchive64Locator:
Ronald Oussoren143cefb2006-06-15 08:14:18 +0000163 return endrec
164
165 if diskno != 0 or disks != 1:
166 raise BadZipfile("zipfiles that span multiple disks are not supported")
167
Tim Petersa608bb22006-06-15 18:06:29 +0000168 # Assume no 'zip64 extensible data'
Martin v. Löwis8c436412008-07-03 12:51:14 +0000169 fpin.seek(offset - sizeEndCentDir64Locator - sizeEndCentDir64, 2)
170 data = fpin.read(sizeEndCentDir64)
Ronald Oussoren143cefb2006-06-15 08:14:18 +0000171 sig, sz, create_version, read_version, disk_num, disk_dir, \
172 dircount, dircount2, dirsize, diroffset = \
Amaury Forgeot d'Arcae6d2b92008-07-11 21:28:25 +0000173 struct.unpack(structEndArchive64, data)
174 if sig != stringEndArchive64:
Ronald Oussoren143cefb2006-06-15 08:14:18 +0000175 return endrec
176
177 # Update the original endrec using data from the ZIP64 record
Antoine Pitrouebcd0ce2008-09-05 23:30:23 +0000178 endrec[_ECD_SIGNATURE] = sig
Martin v. Löwis8c436412008-07-03 12:51:14 +0000179 endrec[_ECD_DISK_NUMBER] = disk_num
180 endrec[_ECD_DISK_START] = disk_dir
181 endrec[_ECD_ENTRIES_THIS_DISK] = dircount
182 endrec[_ECD_ENTRIES_TOTAL] = dircount2
183 endrec[_ECD_SIZE] = dirsize
184 endrec[_ECD_OFFSET] = diroffset
Ronald Oussoren143cefb2006-06-15 08:14:18 +0000185 return endrec
186
187
Martin v. Löwis6f6873b2002-10-13 13:54:50 +0000188def _EndRecData(fpin):
189 """Return data from the "End of Central Directory" record, or None.
190
191 The data is a list of the nine items in the ZIP "End of central dir"
192 record followed by a tenth item, the file seek offset of this record."""
Martin v. Löwis8c436412008-07-03 12:51:14 +0000193
194 # Determine file size
195 fpin.seek(0, 2)
196 filesize = fpin.tell()
197
198 # Check to see if this is ZIP file with no archive comment (the
199 # "end of central directory" structure should be the last item in the
200 # file if this is the case).
Amaury Forgeot d'Arc3e5b0272009-07-28 22:15:30 +0000201 try:
202 fpin.seek(-sizeEndCentDir, 2)
203 except IOError:
204 return None
Martin v. Löwis6f6873b2002-10-13 13:54:50 +0000205 data = fpin.read()
Amaury Forgeot d'Arcae6d2b92008-07-11 21:28:25 +0000206 if data[0:4] == stringEndArchive and data[-2:] == "\000\000":
Martin v. Löwis8c436412008-07-03 12:51:14 +0000207 # the signature is correct and there's no comment, unpack structure
Amaury Forgeot d'Arcae6d2b92008-07-11 21:28:25 +0000208 endrec = struct.unpack(structEndArchive, data)
Martin v. Löwis8c436412008-07-03 12:51:14 +0000209 endrec=list(endrec)
210
211 # Append a blank comment and record start offset
212 endrec.append("")
213 endrec.append(filesize - sizeEndCentDir)
Martin v. Löwis8c436412008-07-03 12:51:14 +0000214
Amaury Forgeot d'Arc2407ac92009-01-17 22:43:50 +0000215 # Try to read the "Zip64 end of central directory" structure
216 return _EndRecData64(fpin, -sizeEndCentDir, endrec)
Martin v. Löwis8c436412008-07-03 12:51:14 +0000217
218 # Either this is not a ZIP file, or it is a ZIP file with an archive
219 # comment. Search the end of the file for the "end of central directory"
220 # record signature. The comment is the last item in the ZIP file and may be
221 # up to 64K long. It is assumed that the "end of central directory" magic
222 # number does not appear in the comment.
223 maxCommentStart = max(filesize - (1 << 16) - sizeEndCentDir, 0)
224 fpin.seek(maxCommentStart, 0)
Martin v. Löwis6f6873b2002-10-13 13:54:50 +0000225 data = fpin.read()
Amaury Forgeot d'Arcae6d2b92008-07-11 21:28:25 +0000226 start = data.rfind(stringEndArchive)
Martin v. Löwis8c436412008-07-03 12:51:14 +0000227 if start >= 0:
228 # found the magic number; attempt to unpack and interpret
229 recData = data[start:start+sizeEndCentDir]
Amaury Forgeot d'Arcae6d2b92008-07-11 21:28:25 +0000230 endrec = list(struct.unpack(structEndArchive, recData))
Martin v. Löwis8c436412008-07-03 12:51:14 +0000231 comment = data[start+sizeEndCentDir:]
232 # check that comment length is correct
233 if endrec[_ECD_COMMENT_SIZE] == len(comment):
Martin v. Löwis6f6873b2002-10-13 13:54:50 +0000234 # Append the archive comment and start offset
235 endrec.append(comment)
Martin v. Löwis8c436412008-07-03 12:51:14 +0000236 endrec.append(maxCommentStart + start)
Amaury Forgeot d'Arc2407ac92009-01-17 22:43:50 +0000237
238 # Try to read the "Zip64 end of central directory" structure
239 return _EndRecData64(fpin, maxCommentStart + start - filesize,
240 endrec)
Martin v. Löwis8c436412008-07-03 12:51:14 +0000241
242 # Unable to find a valid end of central directory structure
243 return
Martin v. Löwis6f6873b2002-10-13 13:54:50 +0000244
Fred Drake484d7352000-10-02 21:14:52 +0000245
Ronald Oussoren143cefb2006-06-15 08:14:18 +0000246class ZipInfo (object):
Fred Drake484d7352000-10-02 21:14:52 +0000247 """Class with attributes describing each file in the ZIP archive."""
248
Ronald Oussoren143cefb2006-06-15 08:14:18 +0000249 __slots__ = (
250 'orig_filename',
251 'filename',
252 'date_time',
253 'compress_type',
254 'comment',
255 'extra',
256 'create_system',
257 'create_version',
258 'extract_version',
259 'reserved',
260 'flag_bits',
261 'volume',
262 'internal_attr',
263 'external_attr',
264 'header_offset',
265 'CRC',
266 'compress_size',
267 'file_size',
Gregory P. Smith0c63fc22008-01-20 01:21:03 +0000268 '_raw_time',
Ronald Oussoren143cefb2006-06-15 08:14:18 +0000269 )
270
Guido van Rossum32abe6f2000-03-31 17:30:02 +0000271 def __init__(self, filename="NoName", date_time=(1980,1,1,0,0,0)):
Greg Ward8e36d282003-06-18 00:53:06 +0000272 self.orig_filename = filename # Original file name in archive
Ronald Oussoren143cefb2006-06-15 08:14:18 +0000273
274 # Terminate the file name at the first null byte. Null bytes in file
275 # names are used as tricks by viruses in archives.
Greg Ward8e36d282003-06-18 00:53:06 +0000276 null_byte = filename.find(chr(0))
277 if null_byte >= 0:
278 filename = filename[0:null_byte]
Ronald Oussoren143cefb2006-06-15 08:14:18 +0000279 # This is used to ensure paths in generated ZIP files always use
280 # forward slashes as the directory separator, as required by the
281 # ZIP format specification.
282 if os.sep != "/" and os.sep in filename:
Greg Ward8e36d282003-06-18 00:53:06 +0000283 filename = filename.replace(os.sep, "/")
Ronald Oussoren143cefb2006-06-15 08:14:18 +0000284
Greg Ward8e36d282003-06-18 00:53:06 +0000285 self.filename = filename # Normalized file name
Tim Peterse1190062001-01-15 03:34:38 +0000286 self.date_time = date_time # year, month, day, hour, min, sec
Guido van Rossum32abe6f2000-03-31 17:30:02 +0000287 # Standard values:
Tim Peterse1190062001-01-15 03:34:38 +0000288 self.compress_type = ZIP_STORED # Type of compression for the file
289 self.comment = "" # Comment for each file
290 self.extra = "" # ZIP extra data
Martin v. Löwis00756902006-02-05 17:09:41 +0000291 if sys.platform == 'win32':
292 self.create_system = 0 # System which created ZIP archive
293 else:
294 # Assume everything else is unix-y
295 self.create_system = 3 # System which created ZIP archive
Tim Peterse1190062001-01-15 03:34:38 +0000296 self.create_version = 20 # Version which created ZIP archive
297 self.extract_version = 20 # Version needed to extract archive
298 self.reserved = 0 # Must be zero
299 self.flag_bits = 0 # ZIP flag bits
300 self.volume = 0 # Volume number of file header
301 self.internal_attr = 0 # Internal attributes
302 self.external_attr = 0 # External file attributes
Guido van Rossum32abe6f2000-03-31 17:30:02 +0000303 # Other attributes are set by class ZipFile:
Tim Peterse1190062001-01-15 03:34:38 +0000304 # header_offset Byte offset to the file header
Tim Peterse1190062001-01-15 03:34:38 +0000305 # CRC CRC-32 of the uncompressed file
306 # compress_size Size of the compressed file
307 # file_size Size of the uncompressed file
Guido van Rossum32abe6f2000-03-31 17:30:02 +0000308
309 def FileHeader(self):
Fred Drake484d7352000-10-02 21:14:52 +0000310 """Return the per-file header as a string."""
Guido van Rossum32abe6f2000-03-31 17:30:02 +0000311 dt = self.date_time
312 dosdate = (dt[0] - 1980) << 9 | dt[1] << 5 | dt[2]
Tim Peters3caca232001-12-06 06:23:26 +0000313 dostime = dt[3] << 11 | dt[4] << 5 | (dt[5] // 2)
Guido van Rossum32abe6f2000-03-31 17:30:02 +0000314 if self.flag_bits & 0x08:
Tim Peterse1190062001-01-15 03:34:38 +0000315 # Set these to zero because we write them after the file data
316 CRC = compress_size = file_size = 0
Guido van Rossum32abe6f2000-03-31 17:30:02 +0000317 else:
Tim Peterse1190062001-01-15 03:34:38 +0000318 CRC = self.CRC
319 compress_size = self.compress_size
320 file_size = self.file_size
Ronald Oussoren143cefb2006-06-15 08:14:18 +0000321
322 extra = self.extra
323
324 if file_size > ZIP64_LIMIT or compress_size > ZIP64_LIMIT:
325 # File is larger than what fits into a 4 byte integer,
326 # fall back to the ZIP64 extension
Gregory P. Smithbf02e3b2008-03-19 03:14:41 +0000327 fmt = '<HHQQ'
Ronald Oussoren143cefb2006-06-15 08:14:18 +0000328 extra = extra + struct.pack(fmt,
329 1, struct.calcsize(fmt)-4, file_size, compress_size)
Martin v. Löwis8c436412008-07-03 12:51:14 +0000330 file_size = 0xffffffff
331 compress_size = 0xffffffff
Ronald Oussoren143cefb2006-06-15 08:14:18 +0000332 self.extract_version = max(45, self.extract_version)
333 self.create_version = max(45, self.extract_version)
334
Martin v. Löwis471617d2008-05-05 17:16:58 +0000335 filename, flag_bits = self._encodeFilenameFlags()
Amaury Forgeot d'Arcae6d2b92008-07-11 21:28:25 +0000336 header = struct.pack(structFileHeader, stringFileHeader,
Martin v. Löwis471617d2008-05-05 17:16:58 +0000337 self.extract_version, self.reserved, flag_bits,
Guido van Rossum32abe6f2000-03-31 17:30:02 +0000338 self.compress_type, dostime, dosdate, CRC,
339 compress_size, file_size,
Martin v. Löwis471617d2008-05-05 17:16:58 +0000340 len(filename), len(extra))
341 return header + filename + extra
342
343 def _encodeFilenameFlags(self):
344 if isinstance(self.filename, unicode):
345 try:
346 return self.filename.encode('ascii'), self.flag_bits
347 except UnicodeEncodeError:
348 return self.filename.encode('utf-8'), self.flag_bits | 0x800
349 else:
350 return self.filename, self.flag_bits
351
352 def _decodeFilename(self):
353 if self.flag_bits & 0x800:
354 return self.filename.decode('utf-8')
355 else:
356 return self.filename
Guido van Rossum32abe6f2000-03-31 17:30:02 +0000357
Ronald Oussoren143cefb2006-06-15 08:14:18 +0000358 def _decodeExtra(self):
359 # Try to decode the extra field.
360 extra = self.extra
361 unpack = struct.unpack
362 while extra:
Gregory P. Smithbf02e3b2008-03-19 03:14:41 +0000363 tp, ln = unpack('<HH', extra[:4])
Ronald Oussoren143cefb2006-06-15 08:14:18 +0000364 if tp == 1:
365 if ln >= 24:
Gregory P. Smithbf02e3b2008-03-19 03:14:41 +0000366 counts = unpack('<QQQ', extra[4:28])
Ronald Oussoren143cefb2006-06-15 08:14:18 +0000367 elif ln == 16:
Gregory P. Smithbf02e3b2008-03-19 03:14:41 +0000368 counts = unpack('<QQ', extra[4:20])
Ronald Oussoren143cefb2006-06-15 08:14:18 +0000369 elif ln == 8:
Gregory P. Smithbf02e3b2008-03-19 03:14:41 +0000370 counts = unpack('<Q', extra[4:12])
Ronald Oussoren143cefb2006-06-15 08:14:18 +0000371 elif ln == 0:
372 counts = ()
373 else:
374 raise RuntimeError, "Corrupt extra field %s"%(ln,)
375
376 idx = 0
377
378 # ZIP64 extension (large files and/or large archives)
Gregory P. Smithbf02e3b2008-03-19 03:14:41 +0000379 if self.file_size in (0xffffffffffffffffL, 0xffffffffL):
Ronald Oussoren143cefb2006-06-15 08:14:18 +0000380 self.file_size = counts[idx]
381 idx += 1
382
Martin v. Löwis8c436412008-07-03 12:51:14 +0000383 if self.compress_size == 0xFFFFFFFFL:
Ronald Oussoren143cefb2006-06-15 08:14:18 +0000384 self.compress_size = counts[idx]
385 idx += 1
386
Martin v. Löwis8c436412008-07-03 12:51:14 +0000387 if self.header_offset == 0xffffffffL:
Ronald Oussoren143cefb2006-06-15 08:14:18 +0000388 old = self.header_offset
389 self.header_offset = counts[idx]
390 idx+=1
391
392 extra = extra[ln+4:]
Tim Petersa608bb22006-06-15 18:06:29 +0000393
Guido van Rossum32abe6f2000-03-31 17:30:02 +0000394
Martin v. Löwisc6d626e2007-02-13 09:49:38 +0000395class _ZipDecrypter:
396 """Class to handle decryption of files stored within a ZIP archive.
397
398 ZIP supports a password-based form of encryption. Even though known
399 plaintext attacks have been found against it, it is still useful
Gregory P. Smithda407232008-01-20 01:32:00 +0000400 to be able to get data out of such a file.
Martin v. Löwisc6d626e2007-02-13 09:49:38 +0000401
402 Usage:
403 zd = _ZipDecrypter(mypwd)
404 plain_char = zd(cypher_char)
405 plain_text = map(zd, cypher_text)
406 """
407
408 def _GenerateCRCTable():
409 """Generate a CRC-32 table.
410
411 ZIP encryption uses the CRC32 one-byte primitive for scrambling some
412 internal keys. We noticed that a direct implementation is faster than
413 relying on binascii.crc32().
414 """
415 poly = 0xedb88320
416 table = [0] * 256
417 for i in range(256):
418 crc = i
419 for j in range(8):
420 if crc & 1:
421 crc = ((crc >> 1) & 0x7FFFFFFF) ^ poly
422 else:
423 crc = ((crc >> 1) & 0x7FFFFFFF)
424 table[i] = crc
425 return table
426 crctable = _GenerateCRCTable()
427
428 def _crc32(self, ch, crc):
429 """Compute the CRC32 primitive on one byte."""
430 return ((crc >> 8) & 0xffffff) ^ self.crctable[(crc ^ ord(ch)) & 0xff]
431
432 def __init__(self, pwd):
433 self.key0 = 305419896
434 self.key1 = 591751049
435 self.key2 = 878082192
436 for p in pwd:
437 self._UpdateKeys(p)
438
439 def _UpdateKeys(self, c):
440 self.key0 = self._crc32(c, self.key0)
441 self.key1 = (self.key1 + (self.key0 & 255)) & 4294967295
442 self.key1 = (self.key1 * 134775813 + 1) & 4294967295
443 self.key2 = self._crc32(chr((self.key1 >> 24) & 255), self.key2)
444
445 def __call__(self, c):
446 """Decrypt a single character."""
447 c = ord(c)
448 k = self.key2 | 2
449 c = c ^ (((k * (k^1)) >> 8) & 255)
450 c = chr(c)
451 self._UpdateKeys(c)
452 return c
453
Martin v. Löwis3eb76482007-03-06 10:41:24 +0000454class ZipExtFile:
455 """File-like object for reading an archive member.
Tim Petersea5962f2007-03-12 18:07:52 +0000456 Is returned by ZipFile.open().
Martin v. Löwis3eb76482007-03-06 10:41:24 +0000457 """
Tim Petersea5962f2007-03-12 18:07:52 +0000458
Martin v. Löwis3eb76482007-03-06 10:41:24 +0000459 def __init__(self, fileobj, zipinfo, decrypt=None):
460 self.fileobj = fileobj
461 self.decrypter = decrypt
462 self.bytes_read = 0L
463 self.rawbuffer = ''
464 self.readbuffer = ''
465 self.linebuffer = ''
466 self.eof = False
467 self.univ_newlines = False
468 self.nlSeps = ("\n", )
469 self.lastdiscard = ''
470
471 self.compress_type = zipinfo.compress_type
472 self.compress_size = zipinfo.compress_size
Tim Petersea5962f2007-03-12 18:07:52 +0000473
Martin v. Löwis3eb76482007-03-06 10:41:24 +0000474 self.closed = False
475 self.mode = "r"
476 self.name = zipinfo.filename
477
478 # read from compressed files in 64k blocks
479 self.compreadsize = 64*1024
480 if self.compress_type == ZIP_DEFLATED:
481 self.dc = zlib.decompressobj(-15)
482
483 def set_univ_newlines(self, univ_newlines):
484 self.univ_newlines = univ_newlines
Tim Petersea5962f2007-03-12 18:07:52 +0000485
Martin v. Löwis3eb76482007-03-06 10:41:24 +0000486 # pick line separator char(s) based on universal newlines flag
487 self.nlSeps = ("\n", )
488 if self.univ_newlines:
489 self.nlSeps = ("\r\n", "\r", "\n")
490
491 def __iter__(self):
492 return self
Tim Petersea5962f2007-03-12 18:07:52 +0000493
Martin v. Löwis3eb76482007-03-06 10:41:24 +0000494 def next(self):
495 nextline = self.readline()
496 if not nextline:
497 raise StopIteration()
498
499 return nextline
500
501 def close(self):
502 self.closed = True
503
504 def _checkfornewline(self):
505 nl, nllen = -1, -1
506 if self.linebuffer:
507 # ugly check for cases where half of an \r\n pair was
508 # read on the last pass, and the \r was discarded. In this
509 # case we just throw away the \n at the start of the buffer.
510 if (self.lastdiscard, self.linebuffer[0]) == ('\r','\n'):
511 self.linebuffer = self.linebuffer[1:]
512
Tim Petersea5962f2007-03-12 18:07:52 +0000513 for sep in self.nlSeps:
Martin v. Löwis3eb76482007-03-06 10:41:24 +0000514 nl = self.linebuffer.find(sep)
515 if nl >= 0:
516 nllen = len(sep)
517 return nl, nllen
518
519 return nl, nllen
Tim Petersea5962f2007-03-12 18:07:52 +0000520
Martin v. Löwis3eb76482007-03-06 10:41:24 +0000521 def readline(self, size = -1):
522 """Read a line with approx. size. If size is negative,
Tim Petersea5962f2007-03-12 18:07:52 +0000523 read a whole line.
Martin v. Löwis3eb76482007-03-06 10:41:24 +0000524 """
525 if size < 0:
526 size = sys.maxint
527 elif size == 0:
528 return ''
529
530 # check for a newline already in buffer
531 nl, nllen = self._checkfornewline()
Tim Petersea5962f2007-03-12 18:07:52 +0000532
Martin v. Löwis3eb76482007-03-06 10:41:24 +0000533 if nl >= 0:
534 # the next line was already in the buffer
535 nl = min(nl, size)
536 else:
537 # no line break in buffer - try to read more
538 size -= len(self.linebuffer)
539 while nl < 0 and size > 0:
540 buf = self.read(min(size, 100))
541 if not buf:
542 break
543 self.linebuffer += buf
544 size -= len(buf)
545
546 # check for a newline in buffer
547 nl, nllen = self._checkfornewline()
Tim Petersea5962f2007-03-12 18:07:52 +0000548
Martin v. Löwis3eb76482007-03-06 10:41:24 +0000549 # we either ran out of bytes in the file, or
550 # met the specified size limit without finding a newline,
551 # so return current buffer
552 if nl < 0:
553 s = self.linebuffer
554 self.linebuffer = ''
555 return s
556
557 buf = self.linebuffer[:nl]
558 self.lastdiscard = self.linebuffer[nl:nl + nllen]
559 self.linebuffer = self.linebuffer[nl + nllen:]
560
561 # line is always returned with \n as newline char (except possibly
562 # for a final incomplete line in the file, which is handled above).
563 return buf + "\n"
564
565 def readlines(self, sizehint = -1):
566 """Return a list with all (following) lines. The sizehint parameter
567 is ignored in this implementation.
568 """
569 result = []
570 while True:
571 line = self.readline()
572 if not line: break
573 result.append(line)
574 return result
575
576 def read(self, size = None):
577 # act like file() obj and return empty string if size is 0
578 if size == 0:
579 return ''
580
581 # determine read size
582 bytesToRead = self.compress_size - self.bytes_read
583
584 # adjust read size for encrypted files since the first 12 bytes
585 # are for the encryption/password information
586 if self.decrypter is not None:
587 bytesToRead -= 12
588
589 if size is not None and size >= 0:
590 if self.compress_type == ZIP_STORED:
591 lr = len(self.readbuffer)
592 bytesToRead = min(bytesToRead, size - lr)
593 elif self.compress_type == ZIP_DEFLATED:
594 if len(self.readbuffer) > size:
595 # the user has requested fewer bytes than we've already
596 # pulled through the decompressor; don't read any more
597 bytesToRead = 0
598 else:
599 # user will use up the buffer, so read some more
600 lr = len(self.rawbuffer)
601 bytesToRead = min(bytesToRead, self.compreadsize - lr)
602
603 # avoid reading past end of file contents
604 if bytesToRead + self.bytes_read > self.compress_size:
605 bytesToRead = self.compress_size - self.bytes_read
606
607 # try to read from file (if necessary)
608 if bytesToRead > 0:
609 bytes = self.fileobj.read(bytesToRead)
610 self.bytes_read += len(bytes)
611 self.rawbuffer += bytes
612
613 # handle contents of raw buffer
614 if self.rawbuffer:
615 newdata = self.rawbuffer
616 self.rawbuffer = ''
617
618 # decrypt new data if we were given an object to handle that
619 if newdata and self.decrypter is not None:
620 newdata = ''.join(map(self.decrypter, newdata))
621
622 # decompress newly read data if necessary
623 if newdata and self.compress_type == ZIP_DEFLATED:
624 newdata = self.dc.decompress(newdata)
625 self.rawbuffer = self.dc.unconsumed_tail
626 if self.eof and len(self.rawbuffer) == 0:
Tim Petersea5962f2007-03-12 18:07:52 +0000627 # we're out of raw bytes (both from the file and
628 # the local buffer); flush just to make sure the
Martin v. Löwis3eb76482007-03-06 10:41:24 +0000629 # decompressor is done
630 newdata += self.dc.flush()
631 # prevent decompressor from being used again
632 self.dc = None
633
634 self.readbuffer += newdata
635
636
637 # return what the user asked for
638 if size is None or len(self.readbuffer) <= size:
639 bytes = self.readbuffer
640 self.readbuffer = ''
641 else:
642 bytes = self.readbuffer[:size]
643 self.readbuffer = self.readbuffer[size:]
644
645 return bytes
Tim Petersea5962f2007-03-12 18:07:52 +0000646
Martin v. Löwis3eb76482007-03-06 10:41:24 +0000647
Guido van Rossum32abe6f2000-03-31 17:30:02 +0000648class ZipFile:
Tim Petersa19a1682001-03-29 04:36:09 +0000649 """ Class with methods to open, read, write, close, list zip files.
650
Martin v. Löwis8c436412008-07-03 12:51:14 +0000651 z = ZipFile(file, mode="r", compression=ZIP_STORED, allowZip64=False)
Tim Petersa19a1682001-03-29 04:36:09 +0000652
Fred Drake3d9091e2001-03-26 15:49:24 +0000653 file: Either the path to the file, or a file-like object.
654 If it is a path, the file will be opened and closed by ZipFile.
655 mode: The mode can be either read "r", write "w" or append "a".
656 compression: ZIP_STORED (no compression) or ZIP_DEFLATED (requires zlib).
Ronald Oussoren143cefb2006-06-15 08:14:18 +0000657 allowZip64: if True ZipFile will create files with ZIP64 extensions when
658 needed, otherwise it will raise an exception when this would
659 be necessary.
660
Fred Drake3d9091e2001-03-26 15:49:24 +0000661 """
Fred Drake484d7352000-10-02 21:14:52 +0000662
Fred Drake90eac282001-02-28 05:29:34 +0000663 fp = None # Set here since __del__ checks it
664
Ronald Oussoren143cefb2006-06-15 08:14:18 +0000665 def __init__(self, file, mode="r", compression=ZIP_STORED, allowZip64=False):
Fred Drake484d7352000-10-02 21:14:52 +0000666 """Open the ZIP file with mode read "r", write "w" or append "a"."""
Georg Brandl4b3ab6f2007-07-12 09:59:22 +0000667 if mode not in ("r", "w", "a"):
668 raise RuntimeError('ZipFile() requires mode "r", "w", or "a"')
669
Guido van Rossum32abe6f2000-03-31 17:30:02 +0000670 if compression == ZIP_STORED:
671 pass
672 elif compression == ZIP_DEFLATED:
673 if not zlib:
674 raise RuntimeError,\
Fred Drake5db246d2000-09-29 20:44:48 +0000675 "Compression requires the (missing) zlib module"
Guido van Rossum32abe6f2000-03-31 17:30:02 +0000676 else:
677 raise RuntimeError, "That compression method is not supported"
Georg Brandl4b3ab6f2007-07-12 09:59:22 +0000678
679 self._allowZip64 = allowZip64
680 self._didModify = False
Tim Peterse1190062001-01-15 03:34:38 +0000681 self.debug = 0 # Level of printing: 0 through 3
682 self.NameToInfo = {} # Find file info given name
683 self.filelist = [] # List of ZipInfo instances for archive
684 self.compression = compression # Method of compression
Raymond Hettinger2ca7c192005-02-16 09:27:49 +0000685 self.mode = key = mode.replace('b', '')[0]
Martin v. Löwisc6d626e2007-02-13 09:49:38 +0000686 self.pwd = None
Martin v. Löwis8c436412008-07-03 12:51:14 +0000687 self.comment = ''
Tim Petersa19a1682001-03-29 04:36:09 +0000688
Fred Drake3d9091e2001-03-26 15:49:24 +0000689 # Check if we were passed a file-like object
Walter Dörwald65230a22002-06-03 15:58:32 +0000690 if isinstance(file, basestring):
Fred Drake3d9091e2001-03-26 15:49:24 +0000691 self._filePassed = 0
692 self.filename = file
693 modeDict = {'r' : 'rb', 'w': 'wb', 'a' : 'r+b'}
Martin v. Löwis84f6de92007-02-13 10:10:39 +0000694 try:
695 self.fp = open(file, modeDict[mode])
696 except IOError:
697 if mode == 'a':
698 mode = key = 'w'
699 self.fp = open(file, modeDict[mode])
700 else:
701 raise
Fred Drake3d9091e2001-03-26 15:49:24 +0000702 else:
703 self._filePassed = 1
704 self.fp = file
705 self.filename = getattr(file, 'name', None)
Tim Petersa19a1682001-03-29 04:36:09 +0000706
Guido van Rossum32abe6f2000-03-31 17:30:02 +0000707 if key == 'r':
Guido van Rossum32abe6f2000-03-31 17:30:02 +0000708 self._GetContents()
709 elif key == 'w':
Fred Drake3d9091e2001-03-26 15:49:24 +0000710 pass
Guido van Rossum32abe6f2000-03-31 17:30:02 +0000711 elif key == 'a':
Martin v. Löwis6f6873b2002-10-13 13:54:50 +0000712 try: # See if file is a zip file
713 self._RealGetContents()
Guido van Rossum32abe6f2000-03-31 17:30:02 +0000714 # seek to start of directory and overwrite
Martin v. Löwis6f6873b2002-10-13 13:54:50 +0000715 self.fp.seek(self.start_dir, 0)
716 except BadZipfile: # file is not a zip file, just append
717 self.fp.seek(0, 2)
Guido van Rossum32abe6f2000-03-31 17:30:02 +0000718 else:
Tim Peters7d3bad62001-04-04 18:56:49 +0000719 if not self._filePassed:
720 self.fp.close()
721 self.fp = None
Guido van Rossum32abe6f2000-03-31 17:30:02 +0000722 raise RuntimeError, 'Mode must be "r", "w" or "a"'
723
724 def _GetContents(self):
Tim Peters7d3bad62001-04-04 18:56:49 +0000725 """Read the directory, making sure we close the file if the format
726 is bad."""
727 try:
728 self._RealGetContents()
729 except BadZipfile:
730 if not self._filePassed:
731 self.fp.close()
732 self.fp = None
733 raise
734
735 def _RealGetContents(self):
Fred Drake484d7352000-10-02 21:14:52 +0000736 """Read in the table of contents for the ZIP file."""
Guido van Rossum32abe6f2000-03-31 17:30:02 +0000737 fp = self.fp
Martin v. Löwis6f6873b2002-10-13 13:54:50 +0000738 endrec = _EndRecData(fp)
739 if not endrec:
740 raise BadZipfile, "File is not a zip file"
Guido van Rossum32abe6f2000-03-31 17:30:02 +0000741 if self.debug > 1:
742 print endrec
Martin v. Löwis8c436412008-07-03 12:51:14 +0000743 size_cd = endrec[_ECD_SIZE] # bytes in central directory
744 offset_cd = endrec[_ECD_OFFSET] # offset of central directory
745 self.comment = endrec[_ECD_COMMENT] # archive comment
746
Guido van Rossum32abe6f2000-03-31 17:30:02 +0000747 # "concat" is zero, unless zip was concatenated to another file
Martin v. Löwis8c436412008-07-03 12:51:14 +0000748 concat = endrec[_ECD_LOCATION] - size_cd - offset_cd
Antoine Pitrouebcd0ce2008-09-05 23:30:23 +0000749 if endrec[_ECD_SIGNATURE] == stringEndArchive64:
750 # If Zip64 extension structures are present, account for them
Martin v. Löwis8c436412008-07-03 12:51:14 +0000751 concat -= (sizeEndCentDir64 + sizeEndCentDir64Locator)
752
Guido van Rossum32abe6f2000-03-31 17:30:02 +0000753 if self.debug > 2:
Martin v. Löwis8c436412008-07-03 12:51:14 +0000754 inferred = concat + offset_cd
755 print "given, inferred, offset", offset_cd, inferred, concat
Guido van Rossum32abe6f2000-03-31 17:30:02 +0000756 # self.start_dir: Position of start of central directory
757 self.start_dir = offset_cd + concat
758 fp.seek(self.start_dir, 0)
Ronald Oussoren143cefb2006-06-15 08:14:18 +0000759 data = fp.read(size_cd)
760 fp = cStringIO.StringIO(data)
Guido van Rossum32abe6f2000-03-31 17:30:02 +0000761 total = 0
762 while total < size_cd:
Martin v. Löwis8c436412008-07-03 12:51:14 +0000763 centdir = fp.read(sizeCentralDir)
Amaury Forgeot d'Arcae6d2b92008-07-11 21:28:25 +0000764 if centdir[0:4] != stringCentralDir:
Guido van Rossum32abe6f2000-03-31 17:30:02 +0000765 raise BadZipfile, "Bad magic number for central directory"
766 centdir = struct.unpack(structCentralDir, centdir)
767 if self.debug > 2:
768 print centdir
Fred Drake3e038e52001-02-28 17:56:26 +0000769 filename = fp.read(centdir[_CD_FILENAME_LENGTH])
Guido van Rossum32abe6f2000-03-31 17:30:02 +0000770 # Create ZipInfo instance to store file information
771 x = ZipInfo(filename)
Fred Drake3e038e52001-02-28 17:56:26 +0000772 x.extra = fp.read(centdir[_CD_EXTRA_FIELD_LENGTH])
773 x.comment = fp.read(centdir[_CD_COMMENT_LENGTH])
Ronald Oussoren143cefb2006-06-15 08:14:18 +0000774 x.header_offset = centdir[_CD_LOCAL_HEADER_OFFSET]
Guido van Rossum32abe6f2000-03-31 17:30:02 +0000775 (x.create_version, x.create_system, x.extract_version, x.reserved,
776 x.flag_bits, x.compress_type, t, d,
777 x.CRC, x.compress_size, x.file_size) = centdir[1:12]
778 x.volume, x.internal_attr, x.external_attr = centdir[15:18]
779 # Convert date/time code to (year, month, day, hour, min, sec)
Gregory P. Smith0c63fc22008-01-20 01:21:03 +0000780 x._raw_time = t
Guido van Rossum32abe6f2000-03-31 17:30:02 +0000781 x.date_time = ( (d>>9)+1980, (d>>5)&0xF, d&0x1F,
Fred Drake414ca662000-06-13 18:49:53 +0000782 t>>11, (t>>5)&0x3F, (t&0x1F) * 2 )
Ronald Oussoren143cefb2006-06-15 08:14:18 +0000783
784 x._decodeExtra()
785 x.header_offset = x.header_offset + concat
Martin v. Löwis471617d2008-05-05 17:16:58 +0000786 x.filename = x._decodeFilename()
Guido van Rossum32abe6f2000-03-31 17:30:02 +0000787 self.filelist.append(x)
788 self.NameToInfo[x.filename] = x
Martin v. Löwis8c436412008-07-03 12:51:14 +0000789
790 # update total bytes read from central directory
791 total = (total + sizeCentralDir + centdir[_CD_FILENAME_LENGTH]
792 + centdir[_CD_EXTRA_FIELD_LENGTH]
793 + centdir[_CD_COMMENT_LENGTH])
794
Guido van Rossum32abe6f2000-03-31 17:30:02 +0000795 if self.debug > 2:
796 print "total", total
Ronald Oussoren143cefb2006-06-15 08:14:18 +0000797
Guido van Rossum32abe6f2000-03-31 17:30:02 +0000798
799 def namelist(self):
Fred Drake484d7352000-10-02 21:14:52 +0000800 """Return a list of file names in the archive."""
Guido van Rossum32abe6f2000-03-31 17:30:02 +0000801 l = []
802 for data in self.filelist:
803 l.append(data.filename)
804 return l
805
806 def infolist(self):
Fred Drake484d7352000-10-02 21:14:52 +0000807 """Return a list of class ZipInfo instances for files in the
808 archive."""
Guido van Rossum32abe6f2000-03-31 17:30:02 +0000809 return self.filelist
810
811 def printdir(self):
Fred Drake484d7352000-10-02 21:14:52 +0000812 """Print a table of contents for the zip file."""
Guido van Rossum32abe6f2000-03-31 17:30:02 +0000813 print "%-46s %19s %12s" % ("File Name", "Modified ", "Size")
814 for zinfo in self.filelist:
Raymond Hettinger351e1a32008-01-14 22:58:05 +0000815 date = "%d-%02d-%02d %02d:%02d:%02d" % zinfo.date_time[:6]
Guido van Rossum32abe6f2000-03-31 17:30:02 +0000816 print "%-46s %s %12d" % (zinfo.filename, date, zinfo.file_size)
817
818 def testzip(self):
Fred Drake484d7352000-10-02 21:14:52 +0000819 """Read all the files and check the CRC."""
Antoine Pitrouc5342702008-08-17 13:06:29 +0000820 chunk_size = 2 ** 20
Guido van Rossum32abe6f2000-03-31 17:30:02 +0000821 for zinfo in self.filelist:
822 try:
Antoine Pitrouc5342702008-08-17 13:06:29 +0000823 # Read by chunks, to avoid an OverflowError or a
824 # MemoryError with very large embedded files.
825 f = self.open(zinfo.filename, "r")
826 while f.read(chunk_size): # Check CRC-32
827 pass
Raymond Hettingerc0fac962003-06-27 22:25:03 +0000828 except BadZipfile:
Guido van Rossum32abe6f2000-03-31 17:30:02 +0000829 return zinfo.filename
830
831 def getinfo(self, name):
Fred Drake484d7352000-10-02 21:14:52 +0000832 """Return the instance of ZipInfo given 'name'."""
Georg Brandl4b3ab6f2007-07-12 09:59:22 +0000833 info = self.NameToInfo.get(name)
834 if info is None:
835 raise KeyError(
836 'There is no item named %r in the archive' % name)
837
838 return info
Guido van Rossum32abe6f2000-03-31 17:30:02 +0000839
Martin v. Löwisc6d626e2007-02-13 09:49:38 +0000840 def setpassword(self, pwd):
841 """Set default password for encrypted files."""
842 self.pwd = pwd
843
844 def read(self, name, pwd=None):
Fred Drake484d7352000-10-02 21:14:52 +0000845 """Return file bytes (as a string) for name."""
Martin v. Löwis3eb76482007-03-06 10:41:24 +0000846 return self.open(name, "r", pwd).read()
847
848 def open(self, name, mode="r", pwd=None):
849 """Return file-like object for 'name'."""
850 if mode not in ("r", "U", "rU"):
851 raise RuntimeError, 'open() requires mode "r", "U", or "rU"'
Guido van Rossum32abe6f2000-03-31 17:30:02 +0000852 if not self.fp:
853 raise RuntimeError, \
Fred Drake5db246d2000-09-29 20:44:48 +0000854 "Attempt to read ZIP archive that was already closed"
Ronald Oussoren143cefb2006-06-15 08:14:18 +0000855
Tim Petersea5962f2007-03-12 18:07:52 +0000856 # Only open a new file for instances where we were not
Martin v. Löwis3eb76482007-03-06 10:41:24 +0000857 # given a file object in the constructor
858 if self._filePassed:
859 zef_file = self.fp
860 else:
861 zef_file = open(self.filename, 'rb')
862
Georg Brandl112aa502008-05-20 08:25:48 +0000863 # Make sure we have an info object
864 if isinstance(name, ZipInfo):
865 # 'name' is already an info object
866 zinfo = name
867 else:
868 # Get info object for name
869 zinfo = self.getinfo(name)
Martin v. Löwis3eb76482007-03-06 10:41:24 +0000870
871 zef_file.seek(zinfo.header_offset, 0)
Ronald Oussoren143cefb2006-06-15 08:14:18 +0000872
873 # Skip the file header:
Martin v. Löwis8c436412008-07-03 12:51:14 +0000874 fheader = zef_file.read(sizeFileHeader)
Amaury Forgeot d'Arcae6d2b92008-07-11 21:28:25 +0000875 if fheader[0:4] != stringFileHeader:
Ronald Oussoren143cefb2006-06-15 08:14:18 +0000876 raise BadZipfile, "Bad magic number for file header"
877
878 fheader = struct.unpack(structFileHeader, fheader)
Martin v. Löwis3eb76482007-03-06 10:41:24 +0000879 fname = zef_file.read(fheader[_FH_FILENAME_LENGTH])
Ronald Oussoren143cefb2006-06-15 08:14:18 +0000880 if fheader[_FH_EXTRA_FIELD_LENGTH]:
Martin v. Löwis3eb76482007-03-06 10:41:24 +0000881 zef_file.read(fheader[_FH_EXTRA_FIELD_LENGTH])
Ronald Oussoren143cefb2006-06-15 08:14:18 +0000882
883 if fname != zinfo.orig_filename:
884 raise BadZipfile, \
885 'File name in directory "%s" and header "%s" differ.' % (
886 zinfo.orig_filename, fname)
887
Martin v. Löwis3eb76482007-03-06 10:41:24 +0000888 # check for encrypted flag & handle password
889 is_encrypted = zinfo.flag_bits & 0x1
890 zd = None
Martin v. Löwisc6d626e2007-02-13 09:49:38 +0000891 if is_encrypted:
Martin v. Löwis3eb76482007-03-06 10:41:24 +0000892 if not pwd:
893 pwd = self.pwd
894 if not pwd:
895 raise RuntimeError, "File %s is encrypted, " \
896 "password required for extraction" % name
897
Martin v. Löwisc6d626e2007-02-13 09:49:38 +0000898 zd = _ZipDecrypter(pwd)
899 # The first 12 bytes in the cypher stream is an encryption header
900 # used to strengthen the algorithm. The first 11 bytes are
901 # completely random, while the 12th contains the MSB of the CRC,
Gregory P. Smith0c63fc22008-01-20 01:21:03 +0000902 # or the MSB of the file time depending on the header type
Martin v. Löwisc6d626e2007-02-13 09:49:38 +0000903 # and is used to check the correctness of the password.
Martin v. Löwis3eb76482007-03-06 10:41:24 +0000904 bytes = zef_file.read(12)
Martin v. Löwisc6d626e2007-02-13 09:49:38 +0000905 h = map(zd, bytes[0:12])
Gregory P. Smith0c63fc22008-01-20 01:21:03 +0000906 if zinfo.flag_bits & 0x8:
907 # compare against the file type from extended local headers
908 check_byte = (zinfo._raw_time >> 8) & 0xff
909 else:
910 # compare against the CRC otherwise
911 check_byte = (zinfo.CRC >> 24) & 0xff
912 if ord(h[11]) != check_byte:
913 raise RuntimeError("Bad password for file", name)
Martin v. Löwis3eb76482007-03-06 10:41:24 +0000914
915 # build and return a ZipExtFile
916 if zd is None:
917 zef = ZipExtFile(zef_file, zinfo)
Guido van Rossum32abe6f2000-03-31 17:30:02 +0000918 else:
Martin v. Löwis3eb76482007-03-06 10:41:24 +0000919 zef = ZipExtFile(zef_file, zinfo, zd)
920
921 # set universal newlines on ZipExtFile if necessary
922 if "U" in mode:
923 zef.set_univ_newlines(True)
924 return zef
Guido van Rossum32abe6f2000-03-31 17:30:02 +0000925
Georg Brandl62416bc2008-01-07 18:47:44 +0000926 def extract(self, member, path=None, pwd=None):
927 """Extract a member from the archive to the current working directory,
928 using its full name. Its file information is extracted as accurately
929 as possible. `member' may be a filename or a ZipInfo object. You can
930 specify a different directory using `path'.
931 """
932 if not isinstance(member, ZipInfo):
933 member = self.getinfo(member)
934
935 if path is None:
936 path = os.getcwd()
937
938 return self._extract_member(member, path, pwd)
939
940 def extractall(self, path=None, members=None, pwd=None):
941 """Extract all members from the archive to the current working
942 directory. `path' specifies a different directory to extract to.
943 `members' is optional and must be a subset of the list returned
944 by namelist().
945 """
946 if members is None:
947 members = self.namelist()
948
949 for zipinfo in members:
950 self.extract(zipinfo, path, pwd)
951
952 def _extract_member(self, member, targetpath, pwd):
953 """Extract the ZipInfo object 'member' to a physical
954 file on the path targetpath.
955 """
956 # build the destination pathname, replacing
957 # forward slashes to platform specific separators.
Antoine Pitrou97377bf2009-05-04 21:17:17 +0000958 # Strip trailing path separator, unless it represents the root.
959 if (targetpath[-1:] in (os.path.sep, os.path.altsep)
960 and len(os.path.splitdrive(targetpath)[1]) > 1):
Georg Brandl62416bc2008-01-07 18:47:44 +0000961 targetpath = targetpath[:-1]
962
963 # don't include leading "/" from file name if present
Martin v. Löwis0dfcfc82009-01-24 14:00:33 +0000964 if member.filename[0] == '/':
Georg Brandl62416bc2008-01-07 18:47:44 +0000965 targetpath = os.path.join(targetpath, member.filename[1:])
966 else:
967 targetpath = os.path.join(targetpath, member.filename)
968
969 targetpath = os.path.normpath(targetpath)
970
971 # Create all upper directories if necessary.
972 upperdirs = os.path.dirname(targetpath)
973 if upperdirs and not os.path.exists(upperdirs):
974 os.makedirs(upperdirs)
975
Martin v. Löwis0dfcfc82009-01-24 14:00:33 +0000976 if member.filename[-1] == '/':
Martin v. Löwis0b09c422009-05-24 19:30:52 +0000977 if not os.path.isdir(targetpath):
978 os.mkdir(targetpath)
Martin v. Löwis0dfcfc82009-01-24 14:00:33 +0000979 return targetpath
980
Georg Brandl112aa502008-05-20 08:25:48 +0000981 source = self.open(member, pwd=pwd)
Georg Brandl62416bc2008-01-07 18:47:44 +0000982 target = file(targetpath, "wb")
983 shutil.copyfileobj(source, target)
984 source.close()
985 target.close()
986
987 return targetpath
988
Guido van Rossum32abe6f2000-03-31 17:30:02 +0000989 def _writecheck(self, zinfo):
Fred Drake484d7352000-10-02 21:14:52 +0000990 """Check for errors before writing a file to the archive."""
Raymond Hettinger54f02222002-06-01 14:18:47 +0000991 if zinfo.filename in self.NameToInfo:
Tim Peterse1190062001-01-15 03:34:38 +0000992 if self.debug: # Warning for duplicate names
Guido van Rossum32abe6f2000-03-31 17:30:02 +0000993 print "Duplicate name:", zinfo.filename
994 if self.mode not in ("w", "a"):
995 raise RuntimeError, 'write() requires mode "w" or "a"'
996 if not self.fp:
997 raise RuntimeError, \
Fred Drake5db246d2000-09-29 20:44:48 +0000998 "Attempt to write ZIP archive that was already closed"
Guido van Rossum32abe6f2000-03-31 17:30:02 +0000999 if zinfo.compress_type == ZIP_DEFLATED and not zlib:
1000 raise RuntimeError, \
Fred Drake5db246d2000-09-29 20:44:48 +00001001 "Compression requires the (missing) zlib module"
Guido van Rossum32abe6f2000-03-31 17:30:02 +00001002 if zinfo.compress_type not in (ZIP_STORED, ZIP_DEFLATED):
1003 raise RuntimeError, \
Fred Drake5db246d2000-09-29 20:44:48 +00001004 "That compression method is not supported"
Ronald Oussoren143cefb2006-06-15 08:14:18 +00001005 if zinfo.file_size > ZIP64_LIMIT:
1006 if not self._allowZip64:
1007 raise LargeZipFile("Filesize would require ZIP64 extensions")
1008 if zinfo.header_offset > ZIP64_LIMIT:
1009 if not self._allowZip64:
1010 raise LargeZipFile("Zipfile size would require ZIP64 extensions")
Guido van Rossum32abe6f2000-03-31 17:30:02 +00001011
1012 def write(self, filename, arcname=None, compress_type=None):
Fred Drake484d7352000-10-02 21:14:52 +00001013 """Put the bytes from filename into the archive under the name
1014 arcname."""
Georg Brandl4b3ab6f2007-07-12 09:59:22 +00001015 if not self.fp:
1016 raise RuntimeError(
1017 "Attempt to write to ZIP archive that was already closed")
1018
Guido van Rossum32abe6f2000-03-31 17:30:02 +00001019 st = os.stat(filename)
Martin v. Löwis0dfcfc82009-01-24 14:00:33 +00001020 isdir = stat.S_ISDIR(st.st_mode)
Raymond Hettinger32200ae2002-06-01 19:51:15 +00001021 mtime = time.localtime(st.st_mtime)
Guido van Rossum32abe6f2000-03-31 17:30:02 +00001022 date_time = mtime[0:6]
1023 # Create ZipInfo instance to store file information
1024 if arcname is None:
Georg Brandl8f7c54e2006-02-20 08:40:38 +00001025 arcname = filename
1026 arcname = os.path.normpath(os.path.splitdrive(arcname)[1])
1027 while arcname[0] in (os.sep, os.altsep):
1028 arcname = arcname[1:]
Martin v. Löwis0dfcfc82009-01-24 14:00:33 +00001029 if isdir:
1030 arcname += '/'
Georg Brandl8f7c54e2006-02-20 08:40:38 +00001031 zinfo = ZipInfo(arcname, date_time)
Andrew M. Kuchling55430212004-07-10 15:40:29 +00001032 zinfo.external_attr = (st[0] & 0xFFFF) << 16L # Unix attributes
Guido van Rossum32abe6f2000-03-31 17:30:02 +00001033 if compress_type is None:
Tim Peterse1190062001-01-15 03:34:38 +00001034 zinfo.compress_type = self.compression
Guido van Rossum32abe6f2000-03-31 17:30:02 +00001035 else:
Tim Peterse1190062001-01-15 03:34:38 +00001036 zinfo.compress_type = compress_type
Ronald Oussoren143cefb2006-06-15 08:14:18 +00001037
1038 zinfo.file_size = st.st_size
Finn Bock03a3bb82001-09-05 18:40:33 +00001039 zinfo.flag_bits = 0x00
Tim Peterse1190062001-01-15 03:34:38 +00001040 zinfo.header_offset = self.fp.tell() # Start of header bytes
Ronald Oussoren143cefb2006-06-15 08:14:18 +00001041
1042 self._writecheck(zinfo)
1043 self._didModify = True
Martin v. Löwis0dfcfc82009-01-24 14:00:33 +00001044
1045 if isdir:
1046 zinfo.file_size = 0
1047 zinfo.compress_size = 0
1048 zinfo.CRC = 0
1049 self.filelist.append(zinfo)
1050 self.NameToInfo[zinfo.filename] = zinfo
1051 self.fp.write(zinfo.FileHeader())
1052 return
1053
Benjamin Petersonb91e8ed2009-05-10 02:29:00 +00001054 with open(filename, "rb") as fp:
1055 # Must overwrite CRC and sizes with correct data later
1056 zinfo.CRC = CRC = 0
1057 zinfo.compress_size = compress_size = 0
1058 zinfo.file_size = file_size = 0
1059 self.fp.write(zinfo.FileHeader())
1060 if zinfo.compress_type == ZIP_DEFLATED:
1061 cmpr = zlib.compressobj(zlib.Z_DEFAULT_COMPRESSION,
1062 zlib.DEFLATED, -15)
1063 else:
1064 cmpr = None
1065 while 1:
1066 buf = fp.read(1024 * 8)
1067 if not buf:
1068 break
1069 file_size = file_size + len(buf)
1070 CRC = crc32(buf, CRC) & 0xffffffff
1071 if cmpr:
1072 buf = cmpr.compress(buf)
1073 compress_size = compress_size + len(buf)
1074 self.fp.write(buf)
Guido van Rossum32abe6f2000-03-31 17:30:02 +00001075 if cmpr:
1076 buf = cmpr.flush()
1077 compress_size = compress_size + len(buf)
1078 self.fp.write(buf)
1079 zinfo.compress_size = compress_size
1080 else:
1081 zinfo.compress_size = file_size
1082 zinfo.CRC = CRC
1083 zinfo.file_size = file_size
Finn Bock03a3bb82001-09-05 18:40:33 +00001084 # Seek backwards and write CRC and file sizes
Tim Petersb64bec32001-09-18 02:26:39 +00001085 position = self.fp.tell() # Preserve current position in file
Finn Bock03a3bb82001-09-05 18:40:33 +00001086 self.fp.seek(zinfo.header_offset + 14, 0)
Gregory P. Smithbf02e3b2008-03-19 03:14:41 +00001087 self.fp.write(struct.pack("<LLL", zinfo.CRC, zinfo.compress_size,
Guido van Rossum32abe6f2000-03-31 17:30:02 +00001088 zinfo.file_size))
Finn Bock03a3bb82001-09-05 18:40:33 +00001089 self.fp.seek(position, 0)
Guido van Rossum32abe6f2000-03-31 17:30:02 +00001090 self.filelist.append(zinfo)
1091 self.NameToInfo[zinfo.filename] = zinfo
1092
Just van Rossumb083cb32002-12-12 12:23:32 +00001093 def writestr(self, zinfo_or_arcname, bytes):
Fred Drake484d7352000-10-02 21:14:52 +00001094 """Write a file into the archive. The contents is the string
Just van Rossumb083cb32002-12-12 12:23:32 +00001095 'bytes'. 'zinfo_or_arcname' is either a ZipInfo instance or
1096 the name of the file in the archive."""
1097 if not isinstance(zinfo_or_arcname, ZipInfo):
1098 zinfo = ZipInfo(filename=zinfo_or_arcname,
Raymond Hettinger351e1a32008-01-14 22:58:05 +00001099 date_time=time.localtime(time.time())[:6])
Just van Rossumb083cb32002-12-12 12:23:32 +00001100 zinfo.compress_type = self.compression
Antoine Pitrou5fdfa3e2008-07-25 19:42:26 +00001101 zinfo.external_attr = 0600 << 16
Just van Rossumb083cb32002-12-12 12:23:32 +00001102 else:
1103 zinfo = zinfo_or_arcname
Georg Brandl4b3ab6f2007-07-12 09:59:22 +00001104
1105 if not self.fp:
1106 raise RuntimeError(
1107 "Attempt to write to ZIP archive that was already closed")
1108
Tim Peterse1190062001-01-15 03:34:38 +00001109 zinfo.file_size = len(bytes) # Uncompressed size
Ronald Oussoren143cefb2006-06-15 08:14:18 +00001110 zinfo.header_offset = self.fp.tell() # Start of header bytes
1111 self._writecheck(zinfo)
1112 self._didModify = True
Gregory P. Smithbf02e3b2008-03-19 03:14:41 +00001113 zinfo.CRC = crc32(bytes) & 0xffffffff # CRC-32 checksum
Guido van Rossum32abe6f2000-03-31 17:30:02 +00001114 if zinfo.compress_type == ZIP_DEFLATED:
1115 co = zlib.compressobj(zlib.Z_DEFAULT_COMPRESSION,
1116 zlib.DEFLATED, -15)
1117 bytes = co.compress(bytes) + co.flush()
Tim Peterse1190062001-01-15 03:34:38 +00001118 zinfo.compress_size = len(bytes) # Compressed size
Guido van Rossum32abe6f2000-03-31 17:30:02 +00001119 else:
1120 zinfo.compress_size = zinfo.file_size
Tim Peterse1190062001-01-15 03:34:38 +00001121 zinfo.header_offset = self.fp.tell() # Start of header bytes
Guido van Rossum32abe6f2000-03-31 17:30:02 +00001122 self.fp.write(zinfo.FileHeader())
Guido van Rossum32abe6f2000-03-31 17:30:02 +00001123 self.fp.write(bytes)
Ronald Oussoren143cefb2006-06-15 08:14:18 +00001124 self.fp.flush()
Guido van Rossum32abe6f2000-03-31 17:30:02 +00001125 if zinfo.flag_bits & 0x08:
Tim Peterse1190062001-01-15 03:34:38 +00001126 # Write CRC and file sizes after the file data
Gregory P. Smith26627332009-06-26 07:50:21 +00001127 self.fp.write(struct.pack("<LLL", zinfo.CRC, zinfo.compress_size,
Tim Peterse1190062001-01-15 03:34:38 +00001128 zinfo.file_size))
Guido van Rossum32abe6f2000-03-31 17:30:02 +00001129 self.filelist.append(zinfo)
1130 self.NameToInfo[zinfo.filename] = zinfo
1131
1132 def __del__(self):
Fred Drake484d7352000-10-02 21:14:52 +00001133 """Call the "close()" method in case the user forgot."""
Tim Petersd15f8bb2001-11-28 23:16:40 +00001134 self.close()
Guido van Rossum32abe6f2000-03-31 17:30:02 +00001135
1136 def close(self):
Fred Drake484d7352000-10-02 21:14:52 +00001137 """Close the file, and for mode "w" and "a" write the ending
1138 records."""
Tim Petersd15f8bb2001-11-28 23:16:40 +00001139 if self.fp is None:
1140 return
Tim Petersa608bb22006-06-15 18:06:29 +00001141
Ronald Oussoren143cefb2006-06-15 08:14:18 +00001142 if self.mode in ("w", "a") and self._didModify: # write ending records
Guido van Rossum32abe6f2000-03-31 17:30:02 +00001143 count = 0
1144 pos1 = self.fp.tell()
Tim Peterse1190062001-01-15 03:34:38 +00001145 for zinfo in self.filelist: # write central directory
Guido van Rossum32abe6f2000-03-31 17:30:02 +00001146 count = count + 1
1147 dt = zinfo.date_time
1148 dosdate = (dt[0] - 1980) << 9 | dt[1] << 5 | dt[2]
Tim Peters3caca232001-12-06 06:23:26 +00001149 dostime = dt[3] << 11 | dt[4] << 5 | (dt[5] // 2)
Ronald Oussoren143cefb2006-06-15 08:14:18 +00001150 extra = []
1151 if zinfo.file_size > ZIP64_LIMIT \
1152 or zinfo.compress_size > ZIP64_LIMIT:
1153 extra.append(zinfo.file_size)
1154 extra.append(zinfo.compress_size)
Martin v. Löwis8c436412008-07-03 12:51:14 +00001155 file_size = 0xffffffff
1156 compress_size = 0xffffffff
Ronald Oussoren143cefb2006-06-15 08:14:18 +00001157 else:
1158 file_size = zinfo.file_size
1159 compress_size = zinfo.compress_size
1160
1161 if zinfo.header_offset > ZIP64_LIMIT:
1162 extra.append(zinfo.header_offset)
Martin v. Löwis8c436412008-07-03 12:51:14 +00001163 header_offset = 0xffffffffL
Ronald Oussoren143cefb2006-06-15 08:14:18 +00001164 else:
1165 header_offset = zinfo.header_offset
1166
1167 extra_data = zinfo.extra
1168 if extra:
1169 # Append a ZIP64 field to the extra's
1170 extra_data = struct.pack(
Gregory P. Smithb89a0962008-03-19 01:46:10 +00001171 '<HH' + 'Q'*len(extra),
Ronald Oussoren143cefb2006-06-15 08:14:18 +00001172 1, 8*len(extra), *extra) + extra_data
Tim Petersa608bb22006-06-15 18:06:29 +00001173
Ronald Oussoren143cefb2006-06-15 08:14:18 +00001174 extract_version = max(45, zinfo.extract_version)
1175 create_version = max(45, zinfo.create_version)
1176 else:
1177 extract_version = zinfo.extract_version
1178 create_version = zinfo.create_version
1179
Gregory P. Smithbf02e3b2008-03-19 03:14:41 +00001180 try:
Martin v. Löwis471617d2008-05-05 17:16:58 +00001181 filename, flag_bits = zinfo._encodeFilenameFlags()
Gregory P. Smithbf02e3b2008-03-19 03:14:41 +00001182 centdir = struct.pack(structCentralDir,
Amaury Forgeot d'Arcae6d2b92008-07-11 21:28:25 +00001183 stringCentralDir, create_version,
Gregory P. Smithbf02e3b2008-03-19 03:14:41 +00001184 zinfo.create_system, extract_version, zinfo.reserved,
Martin v. Löwis471617d2008-05-05 17:16:58 +00001185 flag_bits, zinfo.compress_type, dostime, dosdate,
Gregory P. Smithbf02e3b2008-03-19 03:14:41 +00001186 zinfo.CRC, compress_size, file_size,
Martin v. Löwis471617d2008-05-05 17:16:58 +00001187 len(filename), len(extra_data), len(zinfo.comment),
Gregory P. Smithbf02e3b2008-03-19 03:14:41 +00001188 0, zinfo.internal_attr, zinfo.external_attr,
1189 header_offset)
1190 except DeprecationWarning:
1191 print >>sys.stderr, (structCentralDir,
1192 stringCentralDir, create_version,
1193 zinfo.create_system, extract_version, zinfo.reserved,
1194 zinfo.flag_bits, zinfo.compress_type, dostime, dosdate,
1195 zinfo.CRC, compress_size, file_size,
1196 len(zinfo.filename), len(extra_data), len(zinfo.comment),
1197 0, zinfo.internal_attr, zinfo.external_attr,
1198 header_offset)
1199 raise
Guido van Rossum32abe6f2000-03-31 17:30:02 +00001200 self.fp.write(centdir)
Martin v. Löwis471617d2008-05-05 17:16:58 +00001201 self.fp.write(filename)
Ronald Oussoren143cefb2006-06-15 08:14:18 +00001202 self.fp.write(extra_data)
Guido van Rossum32abe6f2000-03-31 17:30:02 +00001203 self.fp.write(zinfo.comment)
Ronald Oussoren143cefb2006-06-15 08:14:18 +00001204
Guido van Rossum32abe6f2000-03-31 17:30:02 +00001205 pos2 = self.fp.tell()
1206 # Write end-of-zip-archive record
Amaury Forgeot d'Arcd25f87a2009-01-17 16:40:17 +00001207 centDirCount = count
1208 centDirSize = pos2 - pos1
Martin v. Löwis8c436412008-07-03 12:51:14 +00001209 centDirOffset = pos1
Amaury Forgeot d'Arcd25f87a2009-01-17 16:40:17 +00001210 if (centDirCount >= ZIP_FILECOUNT_LIMIT or
1211 centDirOffset > ZIP64_LIMIT or
1212 centDirSize > ZIP64_LIMIT):
Ronald Oussoren143cefb2006-06-15 08:14:18 +00001213 # Need to write the ZIP64 end-of-archive records
1214 zip64endrec = struct.pack(
Amaury Forgeot d'Arcae6d2b92008-07-11 21:28:25 +00001215 structEndArchive64, stringEndArchive64,
Amaury Forgeot d'Arcd25f87a2009-01-17 16:40:17 +00001216 44, 45, 45, 0, 0, centDirCount, centDirCount,
1217 centDirSize, centDirOffset)
Ronald Oussoren143cefb2006-06-15 08:14:18 +00001218 self.fp.write(zip64endrec)
1219
1220 zip64locrec = struct.pack(
Amaury Forgeot d'Arcae6d2b92008-07-11 21:28:25 +00001221 structEndArchive64Locator,
1222 stringEndArchive64Locator, 0, pos2, 1)
Ronald Oussoren143cefb2006-06-15 08:14:18 +00001223 self.fp.write(zip64locrec)
Amaury Forgeot d'Arcd25f87a2009-01-17 16:40:17 +00001224 centDirCount = min(centDirCount, 0xFFFF)
1225 centDirSize = min(centDirSize, 0xFFFFFFFF)
1226 centDirOffset = min(centDirOffset, 0xFFFFFFFF)
Ronald Oussoren143cefb2006-06-15 08:14:18 +00001227
Martin v. Löwis8c436412008-07-03 12:51:14 +00001228 # check for valid comment length
1229 if len(self.comment) >= ZIP_MAX_COMMENT:
1230 if self.debug > 0:
1231 msg = 'Archive comment is too long; truncating to %d bytes' \
1232 % ZIP_MAX_COMMENT
1233 self.comment = self.comment[:ZIP_MAX_COMMENT]
Ronald Oussoren143cefb2006-06-15 08:14:18 +00001234
Amaury Forgeot d'Arcae6d2b92008-07-11 21:28:25 +00001235 endrec = struct.pack(structEndArchive, stringEndArchive,
Amaury Forgeot d'Arcd25f87a2009-01-17 16:40:17 +00001236 0, 0, centDirCount, centDirCount,
1237 centDirSize, centDirOffset, len(self.comment))
Martin v. Löwis8c436412008-07-03 12:51:14 +00001238 self.fp.write(endrec)
1239 self.fp.write(self.comment)
Guido van Rossumf85af612001-04-14 16:45:14 +00001240 self.fp.flush()
Martin v. Löwis8c436412008-07-03 12:51:14 +00001241
Fred Drake3d9091e2001-03-26 15:49:24 +00001242 if not self._filePassed:
1243 self.fp.close()
Guido van Rossum32abe6f2000-03-31 17:30:02 +00001244 self.fp = None
1245
1246
1247class PyZipFile(ZipFile):
Fred Drake484d7352000-10-02 21:14:52 +00001248 """Class to create ZIP archives with Python library files and packages."""
1249
Guido van Rossum32abe6f2000-03-31 17:30:02 +00001250 def writepy(self, pathname, basename = ""):
1251 """Add all files from "pathname" to the ZIP archive.
1252
Fred Drake484d7352000-10-02 21:14:52 +00001253 If pathname is a package directory, search the directory and
1254 all package subdirectories recursively for all *.py and enter
1255 the modules into the archive. If pathname is a plain
1256 directory, listdir *.py and enter all modules. Else, pathname
1257 must be a Python *.py file and the module will be put into the
1258 archive. Added modules are always module.pyo or module.pyc.
1259 This method will compile the module.py into module.pyc if
1260 necessary.
1261 """
Guido van Rossum32abe6f2000-03-31 17:30:02 +00001262 dir, name = os.path.split(pathname)
1263 if os.path.isdir(pathname):
1264 initname = os.path.join(pathname, "__init__.py")
1265 if os.path.isfile(initname):
1266 # This is a package directory, add it
1267 if basename:
1268 basename = "%s/%s" % (basename, name)
1269 else:
1270 basename = name
1271 if self.debug:
1272 print "Adding package in", pathname, "as", basename
1273 fname, arcname = self._get_codename(initname[0:-3], basename)
1274 if self.debug:
1275 print "Adding", arcname
1276 self.write(fname, arcname)
1277 dirlist = os.listdir(pathname)
1278 dirlist.remove("__init__.py")
1279 # Add all *.py files and package subdirectories
1280 for filename in dirlist:
1281 path = os.path.join(pathname, filename)
1282 root, ext = os.path.splitext(filename)
1283 if os.path.isdir(path):
1284 if os.path.isfile(os.path.join(path, "__init__.py")):
1285 # This is a package directory, add it
1286 self.writepy(path, basename) # Recursive call
1287 elif ext == ".py":
1288 fname, arcname = self._get_codename(path[0:-3],
1289 basename)
1290 if self.debug:
1291 print "Adding", arcname
1292 self.write(fname, arcname)
1293 else:
1294 # This is NOT a package directory, add its files at top level
1295 if self.debug:
1296 print "Adding files from directory", pathname
1297 for filename in os.listdir(pathname):
1298 path = os.path.join(pathname, filename)
1299 root, ext = os.path.splitext(filename)
1300 if ext == ".py":
1301 fname, arcname = self._get_codename(path[0:-3],
1302 basename)
1303 if self.debug:
1304 print "Adding", arcname
1305 self.write(fname, arcname)
1306 else:
1307 if pathname[-3:] != ".py":
1308 raise RuntimeError, \
Fred Drake5db246d2000-09-29 20:44:48 +00001309 'Files added with writepy() must end with ".py"'
Guido van Rossum32abe6f2000-03-31 17:30:02 +00001310 fname, arcname = self._get_codename(pathname[0:-3], basename)
1311 if self.debug:
1312 print "Adding file", arcname
1313 self.write(fname, arcname)
1314
1315 def _get_codename(self, pathname, basename):
1316 """Return (filename, archivename) for the path.
1317
Fred Drake484d7352000-10-02 21:14:52 +00001318 Given a module name path, return the correct file path and
1319 archive name, compiling if necessary. For example, given
1320 /python/lib/string, return (/python/lib/string.pyc, string).
1321 """
Guido van Rossum32abe6f2000-03-31 17:30:02 +00001322 file_py = pathname + ".py"
1323 file_pyc = pathname + ".pyc"
1324 file_pyo = pathname + ".pyo"
1325 if os.path.isfile(file_pyo) and \
Raymond Hettinger32200ae2002-06-01 19:51:15 +00001326 os.stat(file_pyo).st_mtime >= os.stat(file_py).st_mtime:
Tim Peterse1190062001-01-15 03:34:38 +00001327 fname = file_pyo # Use .pyo file
Guido van Rossum32abe6f2000-03-31 17:30:02 +00001328 elif not os.path.isfile(file_pyc) or \
Raymond Hettinger32200ae2002-06-01 19:51:15 +00001329 os.stat(file_pyc).st_mtime < os.stat(file_py).st_mtime:
Fred Drake484d7352000-10-02 21:14:52 +00001330 import py_compile
Guido van Rossum32abe6f2000-03-31 17:30:02 +00001331 if self.debug:
1332 print "Compiling", file_py
Martin v. Löwis0c6774d2003-01-15 11:51:06 +00001333 try:
1334 py_compile.compile(file_py, file_pyc, None, True)
1335 except py_compile.PyCompileError,err:
1336 print err.msg
Guido van Rossum32abe6f2000-03-31 17:30:02 +00001337 fname = file_pyc
1338 else:
1339 fname = file_pyc
1340 archivename = os.path.split(fname)[1]
1341 if basename:
1342 archivename = "%s/%s" % (basename, archivename)
1343 return (fname, archivename)
Ronald Oussoren143cefb2006-06-15 08:14:18 +00001344
1345
1346def main(args = None):
1347 import textwrap
1348 USAGE=textwrap.dedent("""\
1349 Usage:
1350 zipfile.py -l zipfile.zip # Show listing of a zipfile
1351 zipfile.py -t zipfile.zip # Test if a zipfile is valid
1352 zipfile.py -e zipfile.zip target # Extract zipfile into target dir
1353 zipfile.py -c zipfile.zip src ... # Create zipfile from sources
1354 """)
1355 if args is None:
1356 args = sys.argv[1:]
1357
1358 if not args or args[0] not in ('-l', '-c', '-e', '-t'):
1359 print USAGE
1360 sys.exit(1)
1361
1362 if args[0] == '-l':
1363 if len(args) != 2:
1364 print USAGE
1365 sys.exit(1)
1366 zf = ZipFile(args[1], 'r')
1367 zf.printdir()
1368 zf.close()
1369
1370 elif args[0] == '-t':
1371 if len(args) != 2:
1372 print USAGE
1373 sys.exit(1)
1374 zf = ZipFile(args[1], 'r')
1375 zf.testzip()
1376 print "Done testing"
1377
1378 elif args[0] == '-e':
1379 if len(args) != 3:
1380 print USAGE
1381 sys.exit(1)
1382
1383 zf = ZipFile(args[1], 'r')
1384 out = args[2]
1385 for path in zf.namelist():
Tim Petersa608bb22006-06-15 18:06:29 +00001386 if path.startswith('./'):
Ronald Oussoren143cefb2006-06-15 08:14:18 +00001387 tgt = os.path.join(out, path[2:])
1388 else:
1389 tgt = os.path.join(out, path)
1390
1391 tgtdir = os.path.dirname(tgt)
1392 if not os.path.exists(tgtdir):
1393 os.makedirs(tgtdir)
Benjamin Petersonb91e8ed2009-05-10 02:29:00 +00001394 with open(tgt, 'wb') as fp:
1395 fp.write(zf.read(path))
Ronald Oussoren143cefb2006-06-15 08:14:18 +00001396 zf.close()
1397
1398 elif args[0] == '-c':
1399 if len(args) < 3:
1400 print USAGE
1401 sys.exit(1)
1402
1403 def addToZip(zf, path, zippath):
1404 if os.path.isfile(path):
1405 zf.write(path, zippath, ZIP_DEFLATED)
1406 elif os.path.isdir(path):
1407 for nm in os.listdir(path):
Tim Petersa608bb22006-06-15 18:06:29 +00001408 addToZip(zf,
Ronald Oussoren143cefb2006-06-15 08:14:18 +00001409 os.path.join(path, nm), os.path.join(zippath, nm))
Tim Petersa608bb22006-06-15 18:06:29 +00001410 # else: ignore
Ronald Oussoren143cefb2006-06-15 08:14:18 +00001411
1412 zf = ZipFile(args[1], 'w', allowZip64=True)
1413 for src in args[2:]:
1414 addToZip(zf, src, os.path.basename(src))
1415
1416 zf.close()
1417
1418if __name__ == "__main__":
1419 main()