Mercurial > touhou
annotate pytouhou/formats/pbg3.py @ 474:ca22df9e70bc
Fix a crash when handling non-ascii paths
author | Thibaut Girka <thib@sitedethib.com> |
---|---|
date | Fri, 27 Sep 2013 18:02:04 +0200 |
parents | 70e2ed71b09c |
children | 1b532e7dd521 |
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 | |
377
70e2ed71b09c
Add meaningful exceptions in format parsing.
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
252
diff
changeset
|
31 from pytouhou.formats import WrongFormatError |
70e2ed71b09c
Add meaningful exceptions in format parsing.
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
252
diff
changeset
|
32 |
58 | 33 logger = get_logger(__name__) |
34 | |
35 | |
0 | 36 class PBG3BitStream(BitStream): |
204
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
37 """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
|
38 |
0 | 39 def read_int(self): |
204
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
40 """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
|
41 |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
42 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
|
43 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
|
44 """ |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
45 |
0 | 46 size = self.read(2) |
47 return self.read((size + 1) * 8) | |
48 | |
49 | |
50 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
|
51 """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
|
52 |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
53 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
|
54 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
|
55 """ |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
56 |
0 | 57 string = [] |
58 for i in range(maxsize): | |
59 byte = self.read(8) | |
60 if byte == 0: | |
61 break | |
62 string.append(byte) | |
63 return ''.join(chr(byte) for byte in string) | |
64 | |
65 | |
66 | |
204
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
67 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
|
68 |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
69 |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
70 |
0 | 71 class PBG3(object): |
204
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
72 """Handle PBG3 archive files. |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
73 |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
74 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
|
75 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
|
76 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
|
77 |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
78 Instance variables: |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
79 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
|
80 bitstream -- PBG3BitStream object |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
81 """ |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
82 |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
83 def __init__(self, entries=None, bitstream=None): |
252
b5c7369abd7c
Improve data reading perfs
Thibaut Girka <thib@sitedethib.com>
parents:
204
diff
changeset
|
84 self.entries = entries or {} |
0 | 85 self.bitstream = bitstream #TODO |
86 | |
87 | |
97 | 88 def __enter__(self): |
89 return self | |
90 | |
91 | |
92 def __exit__(self, type, value, traceback): | |
93 return self.bitstream.__exit__(type, value, traceback) | |
94 | |
95 | |
0 | 96 @classmethod |
97 def read(cls, file): | |
204
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
98 """Read a PBG3 file. |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
99 |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
100 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
|
101 Return a PBG3 instance otherwise. |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
102 """ |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
103 |
0 | 104 magic = file.read(4) |
105 if magic != b'PBG3': | |
377
70e2ed71b09c
Add meaningful exceptions in format parsing.
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
252
diff
changeset
|
106 raise WrongFormatError(magic) |
0 | 107 |
108 bitstream = PBG3BitStream(file) | |
109 entries = {} | |
110 | |
111 nb_entries = bitstream.read_int() | |
112 offset = bitstream.read_int() | |
113 bitstream.seek(offset) | |
114 for i in range(nb_entries): | |
115 unknown1 = bitstream.read_int() | |
116 unknown2 = bitstream.read_int() | |
117 checksum = bitstream.read_int() # Checksum of *compressed data* | |
118 offset = bitstream.read_int() | |
119 size = bitstream.read_int() | |
120 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
|
121 entries[name] = PBG3Entry(unknown1, unknown2, checksum, offset, size) |
0 | 122 |
123 return PBG3(entries, bitstream) | |
124 | |
125 | |
126 def list_files(self): | |
204
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
127 """List files present in the archive.""" |
0 | 128 return self.entries.keys() |
129 | |
130 | |
131 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
|
132 """Extract a given file. |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
133 |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
134 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
|
135 Otherwise, raise an exception. |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
136 |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
137 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
|
138 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
|
139 """ |
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
140 |
0 | 141 unkwn1, unkwn2, checksum, offset, size = self.entries[filename] |
142 self.bitstream.seek(offset) | |
143 data = lzss.decompress(self.bitstream, size) | |
144 if check: | |
204
88361534c77e
Add some documentation (argh, so much left to document!)
Thibaut Girka <thib@sitedethib.com>
parents:
97
diff
changeset
|
145 # Verify the checksum |
0 | 146 compressed_size = self.bitstream.io.tell() - offset |
147 self.bitstream.seek(offset) | |
148 value = 0 | |
149 for c in self.bitstream.io.read(compressed_size): | |
150 value += ord(c) | |
151 value &= 0xFFFFFFFF | |
152 if value != checksum: | |
58 | 153 logger.warn('corrupted data!') |
0 | 154 return data |
155 |