Lib/zipfile.py

   1 "Read and write ZIP files."
   2 # Written by James C. Ahlstrom jim@interet.com
   3 # All rights transferred to CNRI pursuant to the Python contribution agreement
   4
   5 import struct, os, time
   6 import binascii
   7
   8 try:
   9     import zlib # We may need its compression method
  10 except:
  11     zlib = None
  12
  13 class BadZipfile(Exception):
  14     pass
  15 error = BadZipfile      # The exception raised by this module
  16
  17 # constants for Zip file compression methods
  18 ZIP_STORED = 0
  19 ZIP_DEFLATED = 8
  20 # Other ZIP compression methods not supported
  21
  22 # Here are some struct module formats for reading headers
  23 structEndArchive = "<4s4H2lH"     # 9 items, end of archive, 22 bytes
  24 stringEndArchive = "PK\005\006"   # magic number for end of archive record
  25 structCentralDir = "<4s4B4H3l5H2l"# 19 items, central directory, 46 bytes
  26 stringCentralDir = "PK\001\002"   # magic number for central directory
  27 structFileHeader = "<4s2B4H3l2H"  # 12 items, file header record, 30 bytes
  28 stringFileHeader = "PK\003\004"   # magic number for file header
  29
  30
  31 def is_zipfile(filename):
  32     """Quickly see if file is a ZIP file by checking the magic number.
  33
  34     Will not accept a ZIP archive with an ending comment.
  35     """
  36     try:
  37         fpin = open(filename, "rb")
  38         fpin.seek(-22, 2)               # Seek to end-of-file record
  39         endrec = fpin.read()
  40         fpin.close()
  41         if endrec[0:4] == "PK\005\006" and endrec[-2:] == "\000\000":
  42             return 1    # file has correct magic number
  43     except:
  44         pass
  45
  46
  47 class ZipInfo:
  48     """Class with attributes describing each file in the ZIP archive."""
  49
  50     def __init__(self, filename="NoName", date_time=(1980,1,1,0,0,0)):
  51         self.filename = filename        # Name of the file in the archive
  52         self.date_time = date_time      # year, month, day, hour, min, sec
  53         # Standard values:
  54         self.compress_type = ZIP_STORED # Type of compression for the file
  55         self.comment = ""               # Comment for each file
  56         self.extra = ""                 # ZIP extra data
  57         self.create_system = 0          # System which created ZIP archive
  58         self.create_version = 20        # Version which created ZIP archive
  59         self.extract_version = 20       # Version needed to extract archive
  60         self.reserved = 0               # Must be zero
  61         self.flag_bits = 0              # ZIP flag bits
  62         self.volume = 0                 # Volume number of file header
  63         self.internal_attr = 0          # Internal attributes
  64         self.external_attr = 0          # External file attributes
  65         # Other attributes are set by class ZipFile:
  66         # header_offset         Byte offset to the file header
  67         # file_offset           Byte offset to the start of the file data
  68         # CRC                   CRC-32 of the uncompressed file
  69         # compress_size         Size of the compressed file
  70         # file_size             Size of the uncompressed file
  71
  72     def FileHeader(self):
  73         """Return the per-file header as a string."""
  74         dt = self.date_time
  75         dosdate = (dt[0] - 1980) << 9 | dt[1] << 5 | dt[2]
  76         dostime = dt[3] << 11 | dt[4] << 5 | dt[5] / 2
  77         if self.flag_bits & 0x08:
  78           # Set these to zero because we write them after the file data
  79           CRC = compress_size = file_size = 0
  80         else:
  81           CRC = self.CRC
  82           compress_size = self.compress_size
  83           file_size = self.file_size
  84         header = struct.pack(structFileHeader, stringFileHeader,
  85                  self.extract_version, self.reserved, self.flag_bits,
  86                  self.compress_type, dostime, dosdate, CRC,
  87                  compress_size, file_size,
  88                  len(self.filename), len(self.extra))
  89         return header + self.filename + self.extra
  90
  91
  92 class ZipFile:
  93     """Class with methods to open, read, write, close, list zip files."""
  94
  95     def __init__(self, filename, mode="r", compression=ZIP_STORED):
  96         """Open the ZIP file with mode read "r", write "w" or append "a"."""
  97         if compression == ZIP_STORED:
  98             pass
  99         elif compression == ZIP_DEFLATED:
 100             if not zlib:
 101                 raise RuntimeError,\
 102                       "Compression requires the (missing) zlib module"
 103         else:
 104             raise RuntimeError, "That compression method is not supported"
 105         self.debug = 0  # Level of printing: 0 through 3
 106         self.NameToInfo = {}    # Find file info given name
 107         self.filelist = []      # List of ZipInfo instances for archive
 108         self.compression = compression  # Method of compression
 109         self.filename = filename
 110         self.mode = key = mode[0]
 111         if key == 'r':
 112             self.fp = open(filename, "rb")
 113             self._GetContents()
 114         elif key == 'w':
 115             self.fp = open(filename, "wb")
 116         elif key == 'a':
 117             fp = self.fp = open(filename, "r+b")
 118             fp.seek(-22, 2)             # Seek to end-of-file record
 119             endrec = fp.read()
 120             if endrec[0:4] == stringEndArchive and \
 121                        endrec[-2:] == "\000\000":
 122                 self._GetContents()     # file is a zip file
 123                 # seek to start of directory and overwrite
 124                 fp.seek(self.start_dir, 0)
 125             else:               # file is not a zip file, just append
 126                 fp.seek(0, 2)
 127         else:
 128             raise RuntimeError, 'Mode must be "r", "w" or "a"'
 129
 130     def _GetContents(self):
 131         """Read in the table of contents for the ZIP file."""
 132         fp = self.fp
 133         fp.seek(-22, 2)         # Start of end-of-archive record
 134         filesize = fp.tell() + 22       # Get file size
 135         endrec = fp.read(22)    # Archive must not end with a comment!
 136         if endrec[0:4] != stringEndArchive or endrec[-2:] != "\000\000":
 137             raise BadZipfile, "File is not a zip file, or ends with a comment"
 138         endrec = struct.unpack(structEndArchive, endrec)
 139         if self.debug > 1:
 140             print endrec
 141         size_cd = endrec[5]             # bytes in central directory
 142         offset_cd = endrec[6]   # offset of central directory
 143         x = filesize - 22 - size_cd
 144         # "concat" is zero, unless zip was concatenated to another file
 145         concat = x - offset_cd
 146         if self.debug > 2:
 147             print "given, inferred, offset", offset_cd, x, concat
 148         # self.start_dir:  Position of start of central directory
 149         self.start_dir = offset_cd + concat
 150         fp.seek(self.start_dir, 0)
 151         total = 0
 152         while total < size_cd:
 153             centdir = fp.read(46)
 154             total = total + 46
 155             if centdir[0:4] != stringCentralDir:
 156                 raise BadZipfile, "Bad magic number for central directory"
 157             centdir = struct.unpack(structCentralDir, centdir)
 158             if self.debug > 2:
 159                 print centdir
 160             filename = fp.read(centdir[12])
 161             # Create ZipInfo instance to store file information
 162             x = ZipInfo(filename)
 163             x.extra = fp.read(centdir[13])
 164             x.comment = fp.read(centdir[14])
 165             total = total + centdir[12] + centdir[13] + centdir[14]
 166             x.header_offset = centdir[18] + concat
 167             x.file_offset = x.header_offset + 30 + centdir[12] + centdir[13]
 168             (x.create_version, x.create_system, x.extract_version, x.reserved,
 169                 x.flag_bits, x.compress_type, t, d,
 170                 x.CRC, x.compress_size, x.file_size) = centdir[1:12]
 171             x.volume, x.internal_attr, x.external_attr = centdir[15:18]
 172             # Convert date/time code to (year, month, day, hour, min, sec)
 173             x.date_time = ( (d>>9)+1980, (d>>5)&0xF, d&0x1F,
 174                                      t>>11, (t>>5)&0x3F, (t&0x1F) * 2 )
 175             self.filelist.append(x)
 176             self.NameToInfo[x.filename] = x
 177             if self.debug > 2:
 178                 print "total", total
 179         for data in self.filelist:
 180             fp.seek(data.header_offset, 0)
 181             fheader = fp.read(30)
 182             if fheader[0:4] != stringFileHeader:
 183                 raise BadZipfile, "Bad magic number for file header"
 184             fheader = struct.unpack(structFileHeader, fheader)
 185             fname = fp.read(fheader[10])
 186             if fname != data.filename:
 187                 raise RuntimeError, \
 188                       'File name in directory "%s" and header "%s" differ.' % (
 189                           data.filename, fname)
 190
 191     def namelist(self):
 192         """Return a list of file names in the archive."""
 193         l = []
 194         for data in self.filelist:
 195             l.append(data.filename)
 196         return l
 197
 198     def infolist(self):
 199         """Return a list of class ZipInfo instances for files in the
 200         archive."""
 201         return self.filelist
 202
 203     def printdir(self):
 204         """Print a table of contents for the zip file."""
 205         print "%-46s %19s %12s" % ("File Name", "Modified    ", "Size")
 206         for zinfo in self.filelist:
 207             date = "%d-%02d-%02d %02d:%02d:%02d" % zinfo.date_time
 208             print "%-46s %s %12d" % (zinfo.filename, date, zinfo.file_size)
 209
 210     def testzip(self):
 211         """Read all the files and check the CRC."""
 212         for zinfo in self.filelist:
 213             try:
 214                 self.read(zinfo.filename)       # Check CRC-32
 215             except:
 216                 return zinfo.filename
 217
 218     def getinfo(self, name):
 219         """Return the instance of ZipInfo given 'name'."""
 220         return self.NameToInfo[name]
 221
 222     def read(self, name):
 223         """Return file bytes (as a string) for name."""
 224         if self.mode not in ("r", "a"):
 225             raise RuntimeError, 'read() requires mode "r" or "a"'
 226         if not self.fp:
 227             raise RuntimeError, \
 228                   "Attempt to read ZIP archive that was already closed"
 229         zinfo = self.getinfo(name)
 230         filepos = self.fp.tell()
 231         self.fp.seek(zinfo.file_offset, 0)
 232         bytes = self.fp.read(zinfo.compress_size)
 233         self.fp.seek(filepos, 0)
 234         if zinfo.compress_type == ZIP_STORED:
 235             pass
 236         elif zinfo.compress_type == ZIP_DEFLATED:
 237             if not zlib:
 238                 raise RuntimeError, \
 239                       "De-compression requires the (missing) zlib module"
 240             # zlib compress/decompress code by Jeremy Hylton of CNRI
 241             dc = zlib.decompressobj(-15)
 242             bytes = dc.decompress(bytes)
 243             # need to feed in unused pad byte so that zlib won't choke
 244             ex = dc.decompress('Z') + dc.flush()
 245             if ex:
 246                 bytes = bytes + ex
 247         else:
 248             raise BadZipfile, \
 249                   "Unsupported compression method %d for file %s" % \
 250             (zinfo.compress_type, name)
 251         crc = binascii.crc32(bytes)
 252         if crc != zinfo.CRC:
 253             raise BadZipfile, "Bad CRC-32 for file %s" % name
 254         return bytes
 255
 256     def _writecheck(self, zinfo):
 257         """Check for errors before writing a file to the archive."""
 258         if self.NameToInfo.has_key(zinfo.filename):
 259             if self.debug:      # Warning for duplicate names
 260                 print "Duplicate name:", zinfo.filename
 261         if self.mode not in ("w", "a"):
 262             raise RuntimeError, 'write() requires mode "w" or "a"'
 263         if not self.fp:
 264             raise RuntimeError, \
 265                   "Attempt to write ZIP archive that was already closed"
 266         if zinfo.compress_type == ZIP_DEFLATED and not zlib:
 267             raise RuntimeError, \
 268                   "Compression requires the (missing) zlib module"
 269         if zinfo.compress_type not in (ZIP_STORED, ZIP_DEFLATED):
 270             raise RuntimeError, \
 271                   "That compression method is not supported"
 272
 273     def write(self, filename, arcname=None, compress_type=None):
 274         """Put the bytes from filename into the archive under the name
 275         arcname."""
 276         st = os.stat(filename)
 277         mtime = time.localtime(st[8])
 278         date_time = mtime[0:6]
 279         # Create ZipInfo instance to store file information
 280         if arcname is None:
 281           zinfo = ZipInfo(filename, date_time)
 282         else:
 283           zinfo = ZipInfo(arcname, date_time)
 284         zinfo.external_attr = st[0] << 16       # Unix attributes
 285         if compress_type is None:
 286           zinfo.compress_type = self.compression
 287         else:
 288           zinfo.compress_type = compress_type
 289         self._writecheck(zinfo)
 290         fp = open(filename, "rb")
 291         zinfo.flag_bits = 0x08
 292         zinfo.header_offset = self.fp.tell()    # Start of header bytes
 293         self.fp.write(zinfo.FileHeader())
 294         zinfo.file_offset = self.fp.tell()      # Start of file bytes
 295         CRC = 0
 296         compress_size = 0
 297         file_size = 0
 298         if zinfo.compress_type == ZIP_DEFLATED:
 299             cmpr = zlib.compressobj(zlib.Z_DEFAULT_COMPRESSION,
 300                  zlib.DEFLATED, -15)
 301         else:
 302             cmpr = None
 303         while 1:
 304             buf = fp.read(1024 * 8)
 305             if not buf:
 306                 break
 307             file_size = file_size + len(buf)
 308             CRC = binascii.crc32(buf, CRC)
 309             if cmpr:
 310                 buf = cmpr.compress(buf)
 311                 compress_size = compress_size + len(buf)
 312             self.fp.write(buf)
 313         fp.close()
 314         if cmpr:
 315             buf = cmpr.flush()
 316             compress_size = compress_size + len(buf)
 317             self.fp.write(buf)
 318             zinfo.compress_size = compress_size
 319         else:
 320             zinfo.compress_size = file_size
 321         zinfo.CRC = CRC
 322         zinfo.file_size = file_size
 323         # Write CRC and file sizes after the file data
 324         self.fp.write(struct.pack("<lll", zinfo.CRC, zinfo.compress_size,
 325               zinfo.file_size))
 326         self.filelist.append(zinfo)
 327         self.NameToInfo[zinfo.filename] = zinfo
 328
 329     def writestr(self, zinfo, bytes):
 330         """Write a file into the archive.  The contents is the string
 331         'bytes'."""
 332         self._writecheck(zinfo)
 333         zinfo.file_size = len(bytes)            # Uncompressed size
 334         zinfo.CRC = binascii.crc32(bytes)       # CRC-32 checksum
 335         if zinfo.compress_type == ZIP_DEFLATED:
 336             co = zlib.compressobj(zlib.Z_DEFAULT_COMPRESSION,
 337                  zlib.DEFLATED, -15)
 338             bytes = co.compress(bytes) + co.flush()
 339             zinfo.compress_size = len(bytes)    # Compressed size
 340         else:
 341             zinfo.compress_size = zinfo.file_size
 342         zinfo.header_offset = self.fp.tell()    # Start of header bytes
 343         self.fp.write(zinfo.FileHeader())
 344         zinfo.file_offset = self.fp.tell()      # Start of file bytes
 345         self.fp.write(bytes)
 346         if zinfo.flag_bits & 0x08:
 347           # Write CRC and file sizes after the file data
 348           self.fp.write(struct.pack("<lll", zinfo.CRC, zinfo.compress_size,
 349                 zinfo.file_size))
 350         self.filelist.append(zinfo)
 351         self.NameToInfo[zinfo.filename] = zinfo
 352
 353     def __del__(self):
 354         """Call the "close()" method in case the user forgot."""
 355         if self.fp:
 356             self.fp.close()
 357             self.fp = None
 358
 359     def close(self):
 360         """Close the file, and for mode "w" and "a" write the ending
 361         records."""
 362         if self.mode in ("w", "a"):             # write ending records
 363             count = 0
 364             pos1 = self.fp.tell()
 365             for zinfo in self.filelist:         # write central directory
 366                 count = count + 1
 367                 dt = zinfo.date_time
 368                 dosdate = (dt[0] - 1980) << 9 | dt[1] << 5 | dt[2]
 369                 dostime = dt[3] << 11 | dt[4] << 5 | dt[5] / 2
 370                 centdir = struct.pack(structCentralDir,
 371                   stringCentralDir, zinfo.create_version,
 372                   zinfo.create_system, zinfo.extract_version, zinfo.reserved,
 373                   zinfo.flag_bits, zinfo.compress_type, dostime, dosdate,
 374                   zinfo.CRC, zinfo.compress_size, zinfo.file_size,
 375                   len(zinfo.filename), len(zinfo.extra), len(zinfo.comment),
 376                   0, zinfo.internal_attr, zinfo.external_attr,
 377                   zinfo.header_offset)
 378                 self.fp.write(centdir)
 379                 self.fp.write(zinfo.filename)
 380                 self.fp.write(zinfo.extra)
 381                 self.fp.write(zinfo.comment)
 382             pos2 = self.fp.tell()
 383             # Write end-of-zip-archive record
 384             endrec = struct.pack(structEndArchive, stringEndArchive,
 385                      0, 0, count, count, pos2 - pos1, pos1, 0)
 386             self.fp.write(endrec)
 387         self.fp.close()
 388         self.fp = None
 389
 390
 391 class PyZipFile(ZipFile):
 392     """Class to create ZIP archives with Python library files and packages."""
 393
 394     def writepy(self, pathname, basename = ""):
 395         """Add all files from "pathname" to the ZIP archive.
 396
 397         If pathname is a package directory, search the directory and
 398         all package subdirectories recursively for all *.py and enter
 399         the modules into the archive.  If pathname is a plain
 400         directory, listdir *.py and enter all modules.  Else, pathname
 401         must be a Python *.py file and the module will be put into the
 402         archive.  Added modules are always module.pyo or module.pyc.
 403         This method will compile the module.py into module.pyc if
 404         necessary.
 405         """
 406         dir, name = os.path.split(pathname)
 407         if os.path.isdir(pathname):
 408             initname = os.path.join(pathname, "__init__.py")
 409             if os.path.isfile(initname):
 410                 # This is a package directory, add it
 411                 if basename:
 412                     basename = "%s/%s" % (basename, name)
 413                 else:
 414                     basename = name
 415                 if self.debug:
 416                     print "Adding package in", pathname, "as", basename
 417                 fname, arcname = self._get_codename(initname[0:-3], basename)
 418                 if self.debug:
 419                     print "Adding", arcname
 420                 self.write(fname, arcname)
 421                 dirlist = os.listdir(pathname)
 422                 dirlist.remove("__init__.py")
 423                 # Add all *.py files and package subdirectories
 424                 for filename in dirlist:
 425                     path = os.path.join(pathname, filename)
 426                     root, ext = os.path.splitext(filename)
 427                     if os.path.isdir(path):
 428                         if os.path.isfile(os.path.join(path, "__init__.py")):
 429                             # This is a package directory, add it
 430                             self.writepy(path, basename)  # Recursive call
 431                     elif ext == ".py":
 432                         fname, arcname = self._get_codename(path[0:-3],
 433                                          basename)
 434                         if self.debug:
 435                             print "Adding", arcname
 436                         self.write(fname, arcname)
 437             else:
 438                 # This is NOT a package directory, add its files at top level
 439                 if self.debug:
 440                     print "Adding files from directory", pathname
 441                 for filename in os.listdir(pathname):
 442                     path = os.path.join(pathname, filename)
 443                     root, ext = os.path.splitext(filename)
 444                     if ext == ".py":
 445                         fname, arcname = self._get_codename(path[0:-3],
 446                                          basename)
 447                         if self.debug:
 448                             print "Adding", arcname
 449                         self.write(fname, arcname)
 450         else:
 451             if pathname[-3:] != ".py":
 452                 raise RuntimeError, \
 453                       'Files added with writepy() must end with ".py"'
 454             fname, arcname = self._get_codename(pathname[0:-3], basename)
 455             if self.debug:
 456                 print "Adding file", arcname
 457             self.write(fname, arcname)
 458
 459     def _get_codename(self, pathname, basename):
 460         """Return (filename, archivename) for the path.
 461
 462         Given a module name path, return the correct file path and
 463         archive name, compiling if necessary.  For example, given
 464         /python/lib/string, return (/python/lib/string.pyc, string).
 465         """
 466         file_py  = pathname + ".py"
 467         file_pyc = pathname + ".pyc"
 468         file_pyo = pathname + ".pyo"
 469         if os.path.isfile(file_pyo) and \
 470                             os.stat(file_pyo)[8] >= os.stat(file_py)[8]:
 471             fname = file_pyo    # Use .pyo file
 472         elif not os.path.isfile(file_pyc) or \
 473              os.stat(file_pyc)[8] < os.stat(file_py)[8]:
 474             import py_compile
 475             if self.debug:
 476                 print "Compiling", file_py
 477             py_compile.compile(file_py, file_pyc)
 478             fname = file_pyc
 479         else:
 480             fname = file_pyc
 481         archivename = os.path.split(fname)[1]
 482         if basename:
 483             archivename = "%s/%s" % (basename, archivename)
 484         return (fname, archivename)