annotate pytouhou/formats/score.py @ 305:5492472963b0

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