Mercurial > touhou
annotate pytouhou/formats/score.py @ 377:70e2ed71b09c
Add meaningful exceptions in format parsing.
| author | Emmanuel Gil Peyrot <linkmauve@linkmauve.fr> |
|---|---|
| date | Wed, 29 Aug 2012 18:34:28 +0200 |
| parents | 5492472963b0 |
| children | d1f0bb0b7a17 |
| rev | line source |
|---|---|
| 292 | 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 | |
| 305 | 16 from struct import pack, unpack, Struct |
| 292 | 17 from collections import namedtuple |
| 18 from io import BytesIO | |
| 19 | |
|
377
70e2ed71b09c
Add meaningful exceptions in format parsing.
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
305
diff
changeset
|
20 from pytouhou.formats import ChecksumError |
|
70e2ed71b09c
Add meaningful exceptions in format parsing.
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
305
diff
changeset
|
21 |
| 292 | 22 |
| 23 class TH6Score(object): | |
| 24 entry_types = { | |
| 25 b'TH6K': (Struct('<I'), | |
| 26 namedtuple('TH6K', ('unknown',))), | |
| 27 b'HSCR': (Struct('<IIBBB8sx'), | |
| 28 namedtuple('HSCR', ('unknown', 'score', 'character', | |
| 29 'rank', 'stage', 'name'))), | |
| 30 b'PSCR': (Struct('<IIBBBx'), | |
| 31 namedtuple('PSCR', ('unknown', 'score', 'character', | |
| 32 'rank', 'stage'))), | |
| 33 b'CLRD': (Struct('<I5B5BBx'), | |
| 34 namedtuple('CLRD', ('unknown', | |
| 35 'easy', 'normal', 'hard', 'lunatic', | |
| 36 'extra', | |
| 37 'easy_continue', 'normal_continue', | |
| 38 'hard_continue', 'lunatic_continue', | |
| 39 'extra_continue', | |
| 40 'character'))), | |
| 41 b'CATK': (Struct('<I I HH I 34s H HH'), | |
| 42 namedtuple('CATK', ('unknown', 'unknown2', 'num', | |
| 43 'unknown3', 'padding', | |
| 44 'name', 'padding2', | |
| 45 'seen', | |
| 46 'defeated'))), | |
| 47 } | |
| 48 | |
| 49 def __init__(self): | |
| 50 self.key1 = 0 | |
| 51 self.key2 = 0 | |
| 52 self.unknown1 = 0 | |
| 53 self.unknown2 = 16 | |
| 54 self.unknown3 = 0 | |
| 55 self.unknown4 = 0 | |
| 56 self.entries = [] | |
| 57 | |
| 58 | |
| 59 @classmethod | |
| 60 def read(cls, file, decrypt=True, verify=True): | |
| 61 self = cls() | |
| 62 | |
| 63 # Decrypt data | |
| 64 if decrypt: | |
| 65 decrypted_file = BytesIO() | |
| 66 decrypted_file.write(file.read(1)) | |
| 67 key = 0 | |
| 68 for c in file.read(): | |
| 69 encrypted = ord(c) | |
| 70 key = ((key << 3) & 0xFF) | ((key >> 5) & 7) | |
| 71 clear = encrypted ^ key | |
| 72 key += clear | |
| 73 decrypted_file.write(chr(clear)) | |
| 74 file = decrypted_file | |
| 75 | |
| 76 # Read first-part header | |
| 77 file.seek(0) | |
| 78 self.unknown1, self.key1, checksum = unpack('<BBH', file.read(4)) | |
| 79 | |
| 80 # Verify checksum | |
| 81 if verify: | |
| 82 #TODO: is there more to it? | |
|
377
70e2ed71b09c
Add meaningful exceptions in format parsing.
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
305
diff
changeset
|
83 real_sum = sum(ord(c) for c in file.read()) & 0xFFFF |
|
70e2ed71b09c
Add meaningful exceptions in format parsing.
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
305
diff
changeset
|
84 if checksum != real_sum: |
|
70e2ed71b09c
Add meaningful exceptions in format parsing.
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
305
diff
changeset
|
85 raise ChecksumError(checksum, real_sum) |
| 292 | 86 file.seek(4) |
| 87 | |
| 88 # Read second-part header | |
| 89 data = unpack('<HBBIII', file.read(16)) | |
| 90 self.unknown2, self.key2, self.unknown3, offset, self.unknown4, size = data | |
| 91 | |
| 92 #TODO: verify size | |
| 93 | |
| 94 # Read tags | |
| 95 file.seek(offset) | |
| 96 while True: | |
| 97 tag = file.read(4) | |
| 98 if not tag: | |
| 99 break | |
| 100 size, size2 = unpack('<HH', file.read(4)) | |
| 101 assert size == size2 | |
| 102 assert size >= 8 | |
| 103 data = file.read(size-8) | |
| 104 data = cls.entry_types[tag][0].unpack(data) | |
| 105 data = cls.entry_types[tag][1](*data) | |
| 106 self.entries.append((tag, data)) | |
| 107 | |
| 108 return self | |
| 109 | |
| 110 | |
| 111 def write(self, file, encrypt=True): | |
| 112 if encrypt: | |
| 113 clearfile = BytesIO() | |
| 114 else: | |
| 115 clearfile = file | |
| 116 | |
| 117 # Write data | |
| 118 clearfile.seek(20) | |
| 119 for entry in self.entries: | |
| 120 #TODO | |
| 121 tag, data = entry | |
| 122 format = TH6Score.entry_types[tag][0] | |
| 123 clearfile.write(tag) | |
| 124 clearfile.write(pack('<H', format.size + 8) * 2) | |
| 125 clearfile.write(format.pack(*data)) | |
| 126 | |
| 127 # Patch header | |
| 128 size = clearfile.tell() | |
| 129 clearfile.seek(0) | |
| 130 clearfile.write(pack('<BBHHBBIII', | |
| 131 self.unknown1, self.key1, 0, self.unknown2, | |
| 132 self.key2, self.unknown3, 20, self.unknown4, | |
| 133 size)) | |
| 134 | |
| 135 # Patch checksum | |
| 136 clearfile.seek(4) | |
| 137 checksum = sum(ord(c) for c in clearfile.read()) & 0xFFFF | |
| 138 clearfile.seek(2) | |
| 139 clearfile.write(pack('<H', checksum)) | |
| 140 | |
| 141 # Encrypt | |
| 142 if encrypt: | |
| 143 clearfile.seek(0) | |
| 144 file.write(clearfile.read(1)) | |
| 145 key = 0 | |
| 146 for c in clearfile.read(): | |
| 147 clear = ord(c) | |
| 148 key = ((key << 3) & 0xFF) | ((key >> 5) & 7) | |
| 149 encrypted = clear ^ key | |
| 150 key += clear | |
| 151 file.write(chr(encrypted)) | |
| 152 |