1# This Source Code Form is subject to the terms of the Mozilla Public
2# License, v. 2.0. If a copy of the MPL was not distributed with this
3# file, You can obtain one at http://mozilla.org/MPL/2.0/.
4
5import os
6import time
7import zipfile
8
9from mozbuild.util import lock_file
10
11
12class ZipFile(zipfile.ZipFile):
13  """ Class with methods to open, read, write, close, list zip files.
14
15  Subclassing zipfile.ZipFile to allow for overwriting of existing
16  entries, though only for writestr, not for write.
17  """
18  def __init__(self, file, mode="r", compression=zipfile.ZIP_STORED,
19               lock = False):
20    if lock:
21      assert isinstance(file, basestring)
22      self.lockfile = lock_file(file + '.lck')
23    else:
24      self.lockfile = None
25
26    if mode == 'a' and lock:
27      # appending to a file which doesn't exist fails, but we can't check
28      # existence util we hold the lock
29      if (not os.path.isfile(file)) or os.path.getsize(file) == 0:
30        mode = 'w'
31
32    zipfile.ZipFile.__init__(self, file, mode, compression)
33    self._remove = []
34    self.end = self.fp.tell()
35    self.debug = 0
36
37  def writestr(self, zinfo_or_arcname, bytes):
38    """Write contents into the archive.
39
40    The contents is the argument 'bytes',  'zinfo_or_arcname' is either
41    a ZipInfo instance or the name of the file in the archive.
42    This method is overloaded to allow overwriting existing entries.
43    """
44    if not isinstance(zinfo_or_arcname, zipfile.ZipInfo):
45      zinfo = zipfile.ZipInfo(filename=zinfo_or_arcname,
46                              date_time=time.localtime(time.time()))
47      zinfo.compress_type = self.compression
48      # Add some standard UNIX file access permissions (-rw-r--r--).
49      zinfo.external_attr = (0x81a4 & 0xFFFF) << 16L
50    else:
51      zinfo = zinfo_or_arcname
52
53    # Now to the point why we overwrote this in the first place,
54    # remember the entry numbers if we already had this entry.
55    # Optimizations:
56    # If the entry to overwrite is the last one, just reuse that.
57    # If we store uncompressed and the new content has the same size
58    # as the old, reuse the existing entry.
59
60    doSeek = False # store if we need to seek to the eof after overwriting
61    if self.NameToInfo.has_key(zinfo.filename):
62      # Find the last ZipInfo with our name.
63      # Last, because that's catching multiple overwrites
64      i = len(self.filelist)
65      while i > 0:
66        i -= 1
67        if self.filelist[i].filename == zinfo.filename:
68          break
69      zi = self.filelist[i]
70      if ((zinfo.compress_type == zipfile.ZIP_STORED
71           and zi.compress_size == len(bytes))
72          or (i + 1) == len(self.filelist)):
73        # make sure we're allowed to write, otherwise done by writestr below
74        self._writecheck(zi)
75        # overwrite existing entry
76        self.fp.seek(zi.header_offset)
77        if (i + 1) == len(self.filelist):
78          # this is the last item in the file, just truncate
79          self.fp.truncate()
80        else:
81          # we need to move to the end of the file afterwards again
82          doSeek = True
83        # unhook the current zipinfo, the writestr of our superclass
84        # will add a new one
85        self.filelist.pop(i)
86        self.NameToInfo.pop(zinfo.filename)
87      else:
88        # Couldn't optimize, sadly, just remember the old entry for removal
89        self._remove.append(self.filelist.pop(i))
90    zipfile.ZipFile.writestr(self, zinfo, bytes)
91    self.filelist.sort(lambda l, r: cmp(l.header_offset, r.header_offset))
92    if doSeek:
93      self.fp.seek(self.end)
94    self.end = self.fp.tell()
95
96  def close(self):
97    """Close the file, and for mode "w" and "a" write the ending
98    records.
99
100    Overwritten to compact overwritten entries.
101    """
102    if not self._remove:
103      # we don't have anything special to do, let's just call base
104      r = zipfile.ZipFile.close(self)
105      self.lockfile = None
106      return r
107
108    if self.fp.mode != 'r+b':
109      # adjust file mode if we originally just wrote, now we rewrite
110      self.fp.close()
111      self.fp = open(self.filename, 'r+b')
112    all = map(lambda zi: (zi, True), self.filelist) + \
113        map(lambda zi: (zi, False), self._remove)
114    all.sort(lambda l, r: cmp(l[0].header_offset, r[0].header_offset))
115    # empty _remove for multiple closes
116    self._remove = []
117
118    lengths = [all[i+1][0].header_offset - all[i][0].header_offset
119               for i in xrange(len(all)-1)]
120    lengths.append(self.end - all[-1][0].header_offset)
121    to_pos = 0
122    for (zi, keep), length in zip(all, lengths):
123      if not keep:
124        continue
125      oldoff = zi.header_offset
126      # python <= 2.4 has file_offset
127      if hasattr(zi, 'file_offset'):
128        zi.file_offset = zi.file_offset + to_pos - oldoff
129      zi.header_offset = to_pos
130      self.fp.seek(oldoff)
131      content = self.fp.read(length)
132      self.fp.seek(to_pos)
133      self.fp.write(content)
134      to_pos += length
135    self.fp.truncate()
136    zipfile.ZipFile.close(self)
137    self.lockfile = None
138