comparison pytouhou/formats/score.py @ 292:3a81b607f974

Add TH6 score.dat support.
author Thibaut Girka <thib@sitedethib.com>
date Fri, 17 Feb 2012 21:19:57 +0100
parents
children 5492472963b0
comparison
equal deleted inserted replaced
291:f6b8483a990d 292:3a81b607f974
1 # -*- encoding: utf-8 -*-
2 ##
3 ## Copyright (C) 2012 Thibaut Girka <thib@sitedethib.com>
4 ##
5 ## This program is free software; you can redistribute it and/or modify
6 ## it under the terms of the GNU General Public License as published
7 ## by the Free Software Foundation; version 3 only.
8 ##
9 ## This program is distributed in the hope that it will be useful,
10 ## but WITHOUT ANY WARRANTY; without even the implied warranty of
11 ## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
12 ## GNU General Public License for more details.
13 ##
14
15
16 from struct import pack, unpack, calcsize, Struct
17 from collections import namedtuple
18 from io import BytesIO
19
20
21 class TH6Score(object):
22 entry_types = {
23 b'TH6K': (Struct('<I'),
24 namedtuple('TH6K', ('unknown',))),
25 b'HSCR': (Struct('<IIBBB8sx'),
26 namedtuple('HSCR', ('unknown', 'score', 'character',
27 'rank', 'stage', 'name'))),
28 b'PSCR': (Struct('<IIBBBx'),
29 namedtuple('PSCR', ('unknown', 'score', 'character',
30 'rank', 'stage'))),
31 b'CLRD': (Struct('<I5B5BBx'),
32 namedtuple('CLRD', ('unknown',
33 'easy', 'normal', 'hard', 'lunatic',
34 'extra',
35 'easy_continue', 'normal_continue',
36 'hard_continue', 'lunatic_continue',
37 'extra_continue',
38 'character'))),
39 b'CATK': (Struct('<I I HH I 34s H HH'),
40 namedtuple('CATK', ('unknown', 'unknown2', 'num',
41 'unknown3', 'padding',
42 'name', 'padding2',
43 'seen',
44 'defeated'))),
45 }
46
47 def __init__(self):
48 self.key1 = 0
49 self.key2 = 0
50 self.unknown1 = 0
51 self.unknown2 = 16
52 self.unknown3 = 0
53 self.unknown4 = 0
54 self.entries = []
55
56
57 @classmethod
58 def read(cls, file, decrypt=True, verify=True):
59 self = cls()
60
61 # Decrypt data
62 if decrypt:
63 decrypted_file = BytesIO()
64 decrypted_file.write(file.read(1))
65 key = 0
66 for c in file.read():
67 encrypted = ord(c)
68 key = ((key << 3) & 0xFF) | ((key >> 5) & 7)
69 clear = encrypted ^ key
70 key += clear
71 decrypted_file.write(chr(clear))
72 file = decrypted_file
73
74 # Read first-part header
75 file.seek(0)
76 self.unknown1, self.key1, checksum = unpack('<BBH', file.read(4))
77
78 # Verify checksum
79 if verify:
80 #TODO: is there more to it?
81 if checksum != sum(ord(c) for c in file.read()) & 0xFFFF:
82 raise Exception
83 file.seek(4)
84
85 # Read second-part header
86 data = unpack('<HBBIII', file.read(16))
87 self.unknown2, self.key2, self.unknown3, offset, self.unknown4, size = data
88
89 #TODO: verify size
90
91 # Read tags
92 file.seek(offset)
93 while True:
94 tag = file.read(4)
95 if not tag:
96 break
97 size, size2 = unpack('<HH', file.read(4))
98 assert size == size2
99 assert size >= 8
100 data = file.read(size-8)
101 data = cls.entry_types[tag][0].unpack(data)
102 data = cls.entry_types[tag][1](*data)
103 self.entries.append((tag, data))
104
105 return self
106
107
108 def write(self, file, encrypt=True):
109 if encrypt:
110 clearfile = BytesIO()
111 else:
112 clearfile = file
113
114 # Write data
115 clearfile.seek(20)
116 for entry in self.entries:
117 #TODO
118 tag, data = entry
119 format = TH6Score.entry_types[tag][0]
120 clearfile.write(tag)
121 clearfile.write(pack('<H', format.size + 8) * 2)
122 clearfile.write(format.pack(*data))
123
124 # Patch header
125 size = clearfile.tell()
126 clearfile.seek(0)
127 clearfile.write(pack('<BBHHBBIII',
128 self.unknown1, self.key1, 0, self.unknown2,
129 self.key2, self.unknown3, 20, self.unknown4,
130 size))
131
132 # Patch checksum
133 clearfile.seek(4)
134 checksum = sum(ord(c) for c in clearfile.read()) & 0xFFFF
135 clearfile.seek(2)
136 clearfile.write(pack('<H', checksum))
137
138 # Encrypt
139 if encrypt:
140 clearfile.seek(0)
141 file.write(clearfile.read(1))
142 key = 0
143 for c in clearfile.read():
144 clear = ord(c)
145 key = ((key << 3) & 0xFF) | ((key >> 5) & 7)
146 encrypted = clear ^ key
147 key += clear
148 file.write(chr(encrypted))
149