Mercurial > touhou
annotate pytouhou/formats/pbg3.py @ 368:71cd4461bb7f
Minor optimizations
author | Thibaut Girka <thib@sitedethib.com> |
---|---|
date | Wed, 11 Jul 2012 15:19:44 +0200 |
parents | b5c7369abd7c |
children | 70e2ed71b09c |
rev | line source |
---|---|
52
ab826bc29aa2
Add some documentation, GPLv3 headers, README and COPYING file.
Thibaut Girka <thib@sitedethib.com>
parents:
0
diff
changeset
|
1 # -*- encoding: utf-8 -*- |
ab826bc29aa2
Add some documentation, GPLv3 headers, README and COPYING file.
Thibaut Girka <thib@sitedethib.com>
parents:
0
diff
changeset
|
2 ## |
ab826bc29aa2
Add some documentation, GPLv3 headers, README and COPYING file.
Thibaut Girka <thib@sitedethib.com>
parents:
0
diff
changeset
|
3 ## Copyright (C) 2011 Thibaut Girka <thib@sitedethib.com> |
ab826bc29aa2
Add some documentation, GPLv3 headers, README and COPYING file.
Thibaut Girka <thib@sitedethib.com>
parents:
0
diff
changeset
|
4 ## |
ab826bc29aa2
Add some documentation, GPLv3 headers, README and COPYING file.
Thibaut Girka <thib@sitedethib.com>
parents:
0
diff
changeset
|
5 ## This program is free software; you can redistribute it and/or modify |
ab826bc29aa2
Add some documentation, GPLv3 headers, README and COPYING file.
Thibaut Girka <thib@sitedethib.com>
parents:
0
diff
changeset
|
6 ## it under the terms of the GNU General Public License as published |
ab826bc29aa2
Add some documentation, GPLv3 headers, README and COPYING file.
Thibaut Girka <thib@sitedethib.com>
parents:
0
diff
changeset
|
7 ## by the Free Software Foundation; version 3 only. |
ab826bc29aa2
Add some documentation, GPLv3 headers, README and COPYING file.
Thibaut Girka <thib@sitedethib.com>
parents:
0
diff
changeset
|
8 ## |
ab826bc29aa2
Add some documentation, GPLv3 headers, README and COPYING file.
Thibaut Girka <thib@sitedethib.com>
parents:
0
diff
changeset
|
9 ## This program is distributed in the hope that it will be useful, |
ab826bc29aa2
Add some documentation, GPLv3 headers, README and COPYING file.
Thibaut Girka <thib@sitedethib.com>
parents:
0
diff
changeset
|
10 ## but WITHOUT ANY WARRANTY; without even the implied warranty of |
ab826bc29aa2
Add some documentation, GPLv3 headers, README and COPYING file.
Thibaut Girka <thib@sitedethib.com>
parents:
0
diff
changeset
|
11 ## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the |
ab826bc29aa2
Add some documentation, GPLv3 headers, README and COPYING file.
Thibaut Girka <thib@sitedethib.com>
parents:
0
diff
changeset
|
12 ## GNU General Public License for more details. |
ab826bc29aa2
Add some documentation, GPLv3 headers, README and COPYING file.
Thibaut Girka <thib@sitedethib.com>
parents:
0
diff
changeset
|
13 ## |
ab826bc29aa2
Add some documentation, GPLv3 headers, README and COPYING file.
Thibaut Girka <thib@sitedethib.com>
parents:
0
diff
changeset
|
14 |
204
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
15 """PBG3 archive files handling. |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
16 |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
17 This module provides classes for handling the PBG3 file format. |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
18 The PBG3 format is the archive format used by Touhou: EoSD. |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
19 |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
20 PBG3 files are merely a bitstream composed of a header, |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
21 a file table, and LZSS-compressed files. |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
22 """ |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
23 |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
24 from collections import namedtuple |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
25 |
58 | 26 from pytouhou.utils.bitstream import BitStream |
27 import pytouhou.utils.lzss as lzss | |
28 | |
29 from pytouhou.utils.helpers import get_logger | |
30 | |
31 logger = get_logger(__name__) | |
32 | |
33 | |
0 | 34 class PBG3BitStream(BitStream): |
204
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
35 """Helper class to handle strings and integers in PBG3 bitstreams.""" |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
36 |
0 | 37 def read_int(self): |
204
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
38 """Read an integer from the bitstream. |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
39 |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
40 Integers have variable sizes. They begin with a two-bit value indicating |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
41 the number of (non-aligned) bytes to read. |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
42 """ |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
43 |
0 | 44 size = self.read(2) |
45 return self.read((size + 1) * 8) | |
46 | |
47 | |
48 def read_string(self, maxsize): | |
204
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
49 """Read a string from the bitstream. |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
50 |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
51 Strings are stored as standard NULL-termianted sequences of bytes. |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
52 The only catch is that they are not byte-aligned. |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
53 """ |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
54 |
0 | 55 string = [] |
56 for i in range(maxsize): | |
57 byte = self.read(8) | |
58 if byte == 0: | |
59 break | |
60 string.append(byte) | |
61 return ''.join(chr(byte) for byte in string) | |
62 | |
63 | |
64 | |
204
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
65 PBG3Entry = namedtuple('PBG3Entry', 'unknown1 unknown2 checksum offset size') |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
66 |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
67 |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
68 |
0 | 69 class PBG3(object): |
204
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
70 """Handle PBG3 archive files. |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
71 |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
72 PBG3 is a file archive format used in Touhou 6: EoSD. |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
73 This class provides a representation of such files, as well as functions to |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
74 read and extract files from a PBG3 archive. |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
75 |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
76 Instance variables: |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
77 entries -- list of PBG3Entry objects describing files present in the archive |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
78 bitstream -- PBG3BitStream object |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
79 """ |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
80 |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
81 def __init__(self, entries=None, bitstream=None): |
252
b5c7369abd7c
Improve data reading perfs
Thibaut Girka <thib@sitedethib.com>
parents:
204
diff
changeset
|
82 self.entries = entries or {} |
0 | 83 self.bitstream = bitstream #TODO |
84 | |
85 | |
97 | 86 def __enter__(self): |
87 return self | |
88 | |
89 | |
90 def __exit__(self, type, value, traceback): | |
91 return self.bitstream.__exit__(type, value, traceback) | |
92 | |
93 | |
0 | 94 @classmethod |
95 def read(cls, file): | |
204
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
96 """Read a PBG3 file. |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
97 |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
98 Raise an exception if the file is invalid. |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
99 Return a PBG3 instance otherwise. |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
100 """ |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
101 |
0 | 102 magic = file.read(4) |
103 if magic != b'PBG3': | |
104 raise Exception #TODO | |
105 | |
106 bitstream = PBG3BitStream(file) | |
107 entries = {} | |
108 | |
109 nb_entries = bitstream.read_int() | |
110 offset = bitstream.read_int() | |
111 bitstream.seek(offset) | |
112 for i in range(nb_entries): | |
113 unknown1 = bitstream.read_int() | |
114 unknown2 = bitstream.read_int() | |
115 checksum = bitstream.read_int() # Checksum of *compressed data* | |
116 offset = bitstream.read_int() | |
117 size = bitstream.read_int() | |
118 name = bitstream.read_string(255).decode('ascii') | |
204
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
119 entries[name] = PBG3Entry(unknown1, unknown2, checksum, offset, size) |
0 | 120 |
121 return PBG3(entries, bitstream) | |
122 | |
123 | |
124 def list_files(self): | |
204
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
125 """List files present in the archive.""" |
0 | 126 return self.entries.keys() |
127 | |
128 | |
129 def extract(self, filename, check=False): | |
204
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
130 """Extract a given file. |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
131 |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
132 If “filename” is in the archive, extract it and return its contents. |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
133 Otherwise, raise an exception. |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
134 |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
135 By default, the checksum of the file won't be verified, |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
136 you can however force the verification using the “check” argument. |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
137 """ |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
138 |
0 | 139 unkwn1, unkwn2, checksum, offset, size = self.entries[filename] |
140 self.bitstream.seek(offset) | |
141 data = lzss.decompress(self.bitstream, size) | |
142 if check: | |
204
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
143 # Verify the checksum |
0 | 144 compressed_size = self.bitstream.io.tell() - offset |
145 self.bitstream.seek(offset) | |
146 value = 0 | |
147 for c in self.bitstream.io.read(compressed_size): | |
148 value += ord(c) | |
149 value &= 0xFFFFFFFF | |
150 if value != checksum: | |
58 | 151 logger.warn('corrupted data!') |
0 | 152 return data |
153 |