annotate pytouhou/formats/pbg3.py @ 31:55973a3f1222

Some more optimization!
author Thibaut Girka <thib@sitedethib.com>
date Sat, 13 Aug 2011 00:02:11 +0200
parents 6b2c7af2384c
children ab826bc29aa2
Ignore whitespace changes - Everywhere: Within whitespace: At end of lines:
rev   line source
0
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
1 from pytouhou.utils.bitstream import BitStream
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
2 import pytouhou.utils.lzss as lzss
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
3
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
4
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
5 class PBG3BitStream(BitStream):
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
6 def read_int(self):
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
7 size = self.read(2)
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
8 return self.read((size + 1) * 8)
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
9
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
10
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
11 def read_string(self, maxsize):
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
12 string = []
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
13 for i in range(maxsize):
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
14 byte = self.read(8)
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
15 if byte == 0:
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
16 break
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
17 string.append(byte)
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
18 return ''.join(chr(byte) for byte in string)
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
19
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
20
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
21
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
22 class PBG3(object):
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
23 def __init__(self, entries, bitstream=None):
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
24 self.entries = entries
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
25 self.bitstream = bitstream #TODO
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
26
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 @classmethod
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
29 def read(cls, file):
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
30 magic = file.read(4)
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
31 if magic != b'PBG3':
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
32 raise Exception #TODO
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
33
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
34 bitstream = PBG3BitStream(file)
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
35 entries = {}
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
36
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
37 nb_entries = bitstream.read_int()
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
38 offset = bitstream.read_int()
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
39 bitstream.seek(offset)
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
40 for i in range(nb_entries):
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
41 unknown1 = bitstream.read_int()
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
42 unknown2 = bitstream.read_int()
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
43 checksum = bitstream.read_int() # Checksum of *compressed data*
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
44 offset = bitstream.read_int()
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
45 size = bitstream.read_int()
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
46 name = bitstream.read_string(255).decode('ascii')
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
47 entries[name] = (unknown1, unknown2, checksum, offset, size)
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
48
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
49 return PBG3(entries, bitstream)
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
50
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
51
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
52 def list_files(self):
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
53 return self.entries.keys()
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
54
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
55
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
56 def extract(self, filename, check=False):
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
57 unkwn1, unkwn2, checksum, offset, size = self.entries[filename]
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
58 self.bitstream.seek(offset)
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
59 data = lzss.decompress(self.bitstream, size)
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
60 if check:
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
61 # Checking the checksum
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
62 compressed_size = self.bitstream.io.tell() - offset
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
63 self.bitstream.seek(offset)
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
64 value = 0
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
65 for c in self.bitstream.io.read(compressed_size):
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
66 value += ord(c)
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
67 value &= 0xFFFFFFFF
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
68 if value != checksum:
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
69 print('Warning: corrupted data') #TODO
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
70 return data
6b2c7af2384c Hello Gensokyo _o/
Thibaut Girka <thib@sitedethib.com>
parents:
diff changeset
71