985a858f1cea8394fad51da7a8c465a8d2eb63c3
[backgammon/xgdatatools.git] / xgzarc.py
1 #
2 #   xgzarc.py - XG ZLib archive module
3 #   Copyright (C) 2013,2014  Michael Petch <mpetch@gnubg.org>
4 #                                          <mpetch@capp-sysware.com>
5 #
6 #   This program is free software: you can redistribute it and/or modify
7 #   it under the terms of the GNU General Public License as published by
8 #   the Free Software Foundation, either version 3 of the License, or
9 #   (at your option) any later version.
10 #
11 #   This program is distributed in the hope that it will be useful,
12 #   but WITHOUT ANY WARRANTY; without even the implied warranty of
13 #   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
14 #   GNU General Public License for more details.
15 #
16 #   You should have received a copy of the GNU General Public License
17 #   along with this program.  If not, see <http://www.gnu.org/licenses/>.
18 #
19 #
20 #   This library is an interpetation of ZLBArchive 1.52 data structures.
21 #   Please see: http://www.delphipages.com/comp/zlibarchive-2104.html
22 #
23
24 from __future__ import with_statement as _with
25 import tempfile as _tempfile
26 import struct as _struct
27 import zlib as _zlib
28 import os as _os
29 import xgutils as _xgutils
30
31
32 class Error(Exception):
33
34     def __init__(self, error):
35         self.value = "Zlib archive: %s" % str(error)
36         self.error = error
37
38     def __str__(self):
39         return repr(self.value)
40
41
42 class ArchiveRecord(dict):
43
44     SIZEOFREC = 36
45
46     def __init__(self, **kw):
47         defaults = {
48             'crc': 0,
49             'filecount': 0,
50             'version': 0,
51             'registrysize': 0,
52             'archivesize': 0,
53             'compressedregistry': False,
54             'reserved': []
55             }
56         super(ArchiveRecord, self).__init__(defaults, **kw)
57
58     def __setattr__(self, key, value):
59         self[key] = value
60
61     def __getattr__(self, key):
62         return self[key]
63
64     def fromstream(self, stream):
65         unpacked_data = _struct.unpack('<llllll12B',
66                                        stream.read(self.SIZEOFREC))
67         self.crc = unpacked_data[0] & 0xffffffff
68         self.filecount = unpacked_data[1]
69         self.version = unpacked_data[2]
70         self.registrysize = unpacked_data[3]
71         self.archivesize = unpacked_data[4]
72         self.compressedregistry = bool(unpacked_data[5])
73         self.reserved = unpacked_data[6:]
74
75
76 class FileRecord(dict):
77
78     SIZEOFREC = 532
79
80     def __init__(self, **kw):
81         defaults = {
82             'name': None,
83             'path': None,
84             'osize': 0,
85             'csize': 0,
86             'start': 0,
87             'crc': 0,
88             'compressed': False,
89             'stored': False,
90             'compressionlevel': 0
91             }
92         super(FileRecord, self).__init__(defaults, **kw)
93
94     def __setattr__(self, key, value):
95         self[key] = value
96
97     def __getattr__(self, key):
98         return self[key]
99
100     def fromstream(self, stream):
101         unpacked_data = _struct.unpack('<256B256BllllBBxx',
102                                        stream.read(self.SIZEOFREC))
103         self.name = _xgutils.delphishortstrtostr(unpacked_data[0:256])
104         self.path = _xgutils.delphishortstrtostr(unpacked_data[256:512])
105         self.osize = unpacked_data[512]
106         self.csize = unpacked_data[513]
107         self.start = unpacked_data[514]
108         self.crc = unpacked_data[515] & 0xffffffff
109         self.compressed = bool(unpacked_data[516] == 0)
110         self.compressionlevel = unpacked_data[517]
111
112     def __str__(self):
113         return str(self.todict())
114
115
116 class ZlibArchive(object):
117     __MAXBUFSIZE = 32768
118     __TMP_PREFIX = 'tmpXGI'
119
120     def __init__(self, stream=None, filename=None):
121         self.arcrec = ArchiveRecord()
122         self.arcregistry = []
123         self.startofarcdata = -1
124         self.endofarcdata = -1
125
126         self.filename = filename
127         self.stream = stream
128         if stream is None:
129             self.stream = open(filename, 'rb')
130
131         self.__getarchiveindex()
132
133     def __extractsegment(self, iscompressed=True, numbytes=None):
134         # Extract a stored segment
135         filename = None
136         stream = []
137
138         try:
139             tmpfd, filename = _tempfile.mkstemp(prefix=self.__TMP_PREFIX)
140             with _os.fdopen(tmpfd, "wb") as tmpfile:
141
142                 if (iscompressed):
143                     # Extract a compressed segment
144                     decomp = _zlib.decompressobj()
145                     buf = self.stream.read(self.__MAXBUFSIZE)
146                     stream = decomp.decompress(buf)
147
148                     if len(stream) <= 0:
149                         raise IOError()
150
151                     tmpfile.write(stream)
152
153                     # Read until we have uncompressed a complete segment
154                     while len(decomp.unused_data) == 0:
155                         block = self.stream.read(self.__MAXBUFSIZE)
156                         if len(block) > 0:
157                             try:
158                                 stream = decomp.decompress(block)
159                                 tmpfile.write(stream)
160                             except:
161                                 break
162                         else:
163                             # EOF reached
164                             break
165
166                 else:
167                     # Extract an uncompressed segment
168                     # Uncompressed segment needs numbytes specified
169                     if numbytes is None:
170                         raise IOError()
171
172                     blksize = self.__MAXBUFSIZE
173                     bytesleft = numbytes
174                     while True:
175                         if bytesleft < blksize:
176                             blksize = bytesleft
177
178                         block = self.stream.read(blksize)
179                         tmpfile.write(block)
180                         bytesleft = bytesleft - blksize
181
182                         if bytesleft == 0:
183                             break
184
185         except (_zlib.error, IOError) as e:
186             _os.unlink(filename)
187             return None
188
189         return filename
190
191     def __getarchiveindex(self):
192
193         try:
194             # Advance to the archive record at the end and retrieve it
195             filerecords = []
196             curstreampos = self.stream.tell()
197
198             self.stream.seek(-ArchiveRecord.SIZEOFREC, _os.SEEK_END)
199             self.endofarcdata = self.stream.tell()
200             self.arcrec.fromstream(self.stream)
201
202             # Position ourselves at the beginning of the archive file index
203             self.stream.seek(-ArchiveRecord.SIZEOFREC -
204                              self.arcrec.registrysize, _os.SEEK_END)
205             self.startofarcdata = self.stream.tell() - self.arcrec.archivesize
206
207             # Compute the CRC32 of all the archive data including file index
208             streamcrc = _xgutils.streamcrc32(
209                     self.stream,
210                     startpos=self.startofarcdata,
211                     numbytes=(self.endofarcdata - self.startofarcdata))
212             if streamcrc != self.arcrec.crc:
213                 raise Error("Archive CRC check failed - file corrupt")
214
215             # Decompress the index into a temporary file
216             idx_filename = self.__extractsegment(
217                     iscompressed=self.arcrec.compressedregistry)
218             if idx_filename is None:
219                 raise Error("Error extracting archive index")
220
221             # Retrieve all the files in the index
222             with open(idx_filename, "rb") as idx_file:
223                 for recordnum in range(0, self.arcrec.filecount):
224                     curidxpos = self.stream.tell()
225
226                     # Retrieve next file index record
227                     filerec = FileRecord()
228                     filerec.fromstream(idx_file)
229                     filerecords.append(filerec)
230
231                     self.stream.seek(curidxpos, 0)
232
233             _os.unlink(idx_filename)
234         finally:
235             self.stream.seek(curstreampos, 0)
236
237         self.arcregistry = filerecords
238
239     def getarchivefile(self, filerec):
240         # Do processing on the temporary file
241         self.stream.seek(filerec.start + self.startofarcdata)
242         tmpfilename = self.__extractsegment(iscompressed=filerec.compressed,
243                                             numbytes=filerec.csize)
244         if tmpfilename is None:
245             raise Error("Error extracting archived file")
246         tmpfile = open(tmpfilename, "rb")
247
248         # Compute the CRC32 on the uncompressed file
249         streamcrc = _xgutils.streamcrc32(tmpfile)
250         if streamcrc != filerec.crc:
251             raise Error("File CRC check failed - file corrupt")
252
253         return tmpfile, tmpfilename
254
255     def setblocksize(self, blksize):
256         self.__MAXBUFSIZE = blksize
257
258
259 if __name__ == '__main__':
260     pass