annotate pytouhou/formats/pbg3.py @ 123:d1c82d43bbf3

Various optimizations
author Thibaut Girka <thib@sitedethib.com>
date Sat, 10 Sep 2011 11:58:24 +0200
parents ac2e5e1c2c3c
children 88361534c77e
Ignore whitespace changes - Everywhere: Within whitespace: At end of lines:
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
58
3da4de9decd0 Use logging module
Thibaut Girka <thib@sitedethib.com>
parents: 52
diff changeset
15 from pytouhou.utils.bitstream import BitStream
3da4de9decd0 Use logging module
Thibaut Girka <thib@sitedethib.com>
parents: 52
diff changeset
16 import pytouhou.utils.lzss as lzss
3da4de9decd0 Use logging module
Thibaut Girka <thib@sitedethib.com>
parents: 52
diff changeset
17
3da4de9decd0 Use logging module
Thibaut Girka <thib@sitedethib.com>
parents: 52
diff changeset
18 from pytouhou.utils.helpers import get_logger
3da4de9decd0 Use logging module
Thibaut Girka <thib@sitedethib.com>
parents: 52
diff changeset
19
3da4de9decd0 Use logging module
Thibaut Girka <thib@sitedethib.com>
parents: 52
diff changeset
20 logger = get_logger(__name__)
3da4de9decd0 Use logging module
Thibaut Girka <thib@sitedethib.com>
parents: 52
diff changeset
21
3da4de9decd0 Use logging module
Thibaut Girka <thib@sitedethib.com>
parents: 52
diff changeset
22
0
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
23 class PBG3BitStream(BitStream):
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
24 def read_int(self):
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
25 size = self.read(2)
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
26 return self.read((size + 1) * 8)
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
27
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
28
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
29 def read_string(self, maxsize):
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
30 string = []
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
31 for i in range(maxsize):
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
32 byte = self.read(8)
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
33 if byte == 0:
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
34 break
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
35 string.append(byte)
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
36 return ''.join(chr(byte) for byte in string)
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
37
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
38
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
39
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
40 class PBG3(object):
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
41 def __init__(self, entries, bitstream=None):
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
42 self.entries = entries
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
43 self.bitstream = bitstream #TODO
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
44
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
45
97
ac2e5e1c2c3c Refactor \o/
Thibaut Girka <thib@sitedethib.com>
parents: 58
diff changeset
46 def __enter__(self):
ac2e5e1c2c3c Refactor \o/
Thibaut Girka <thib@sitedethib.com>
parents: 58
diff changeset
47 return self
ac2e5e1c2c3c Refactor \o/
Thibaut Girka <thib@sitedethib.com>
parents: 58
diff changeset
48
ac2e5e1c2c3c Refactor \o/
Thibaut Girka <thib@sitedethib.com>
parents: 58
diff changeset
49
ac2e5e1c2c3c Refactor \o/
Thibaut Girka <thib@sitedethib.com>
parents: 58
diff changeset
50 def __exit__(self, type, value, traceback):
ac2e5e1c2c3c Refactor \o/
Thibaut Girka <thib@sitedethib.com>
parents: 58
diff changeset
51 return self.bitstream.__exit__(type, value, traceback)
ac2e5e1c2c3c Refactor \o/
Thibaut Girka <thib@sitedethib.com>
parents: 58
diff changeset
52
ac2e5e1c2c3c Refactor \o/
Thibaut Girka <thib@sitedethib.com>
parents: 58
diff changeset
53
0
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
54 @classmethod
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
55 def read(cls, file):
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
56 magic = file.read(4)
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
57 if magic != b'PBG3':
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
58 raise Exception #TODO
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
59
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
60 bitstream = PBG3BitStream(file)
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
61 entries = {}
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
62
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
63 nb_entries = bitstream.read_int()
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
64 offset = bitstream.read_int()
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
65 bitstream.seek(offset)
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
66 for i in range(nb_entries):
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
67 unknown1 = bitstream.read_int()
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
68 unknown2 = bitstream.read_int()
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
69 checksum = bitstream.read_int() # Checksum of *compressed data*
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
70 offset = bitstream.read_int()
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
71 size = bitstream.read_int()
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
72 name = bitstream.read_string(255).decode('ascii')
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
73 entries[name] = (unknown1, unknown2, checksum, offset, size)
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
74
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
75 return PBG3(entries, bitstream)
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
76
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
77
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
78 def list_files(self):
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
79 return self.entries.keys()
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
80
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
81
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
82 def extract(self, filename, check=False):
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
83 unkwn1, unkwn2, checksum, offset, size = self.entries[filename]
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
84 self.bitstream.seek(offset)
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
85 data = lzss.decompress(self.bitstream, size)
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
86 if check:
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
87 # Checking the checksum
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
88 compressed_size = self.bitstream.io.tell() - offset
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
89 self.bitstream.seek(offset)
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
90 value = 0
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
91 for c in self.bitstream.io.read(compressed_size):
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
92 value += ord(c)
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
93 value &= 0xFFFFFFFF
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
94 if value != checksum:
58
3da4de9decd0 Use logging module
Thibaut Girka <thib@sitedethib.com>
parents: 52
diff changeset
95 logger.warn('corrupted data!')
0
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
96 return data
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
97