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