Guest User

Untitled

a guest
May 7th, 2021
50
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. import itertools
  2. import json
  3. import lzma
  4. import zlib
  5. import re
  6. import struct
  7. import sys
  8. import os
  9.  
  10. from collections import Counter
  11.  
  12. # from savegame import TTDSavegameDecoder
  13.  
  14. SCRIPT_DIR = os.path.dirname(os.path.realpath(__file__))
  15.  
  16. def hex_str(s):
  17.     if isinstance(s, (bytes, memoryview)):
  18.         return ':'.join('{:02x}'.format(b) for b in s)
  19.     return ':'.join('{:02x}'.format(ord(c)) for c in s)
  20.  
  21. def fmt_str(s):
  22.     if not s:
  23.         return ''
  24.     try:
  25.         return '%s (%s)' % (s.decode('utf-8'), hex_str(s))
  26.     except Exception as e:
  27.         return hex_str(s)
  28.  
  29.  
  30. def ttd_unpack(fmt, buf, offset=0):
  31.     # print(fmt, len(buf[offset:]), hex_str(buf[offset:]))
  32.     assert isinstance(fmt, str), fmt
  33.     if fmt == '&':
  34.         return [buf[offset:-1]]
  35.     while True:
  36.         pos = fmt.find('z')
  37.         if pos < 0:
  38.             break
  39.         asciiz_start = struct.calcsize(fmt[:pos]) + offset
  40.         asciiz_len, len_size = TTDSavegameDecoder.read_gamma(buf[asciiz_start:])
  41.         fmt = '%s%s%ds%s' % (fmt[:pos], 'x' * len_size, asciiz_len, fmt[pos + 1:])
  42.     while True:
  43.         pos = fmt.find('l')
  44.         if pos < 0:
  45.             break
  46.         list_start = struct.calcsize(fmt[:pos]) + offset
  47.         # print('l', fmt, list_start, buf[list_start:].tobytes())
  48.         list_len = struct.unpack_from('>I', buf, offset=list_start)[0]
  49.         fmt = '%sxxxx%ds%s' % (fmt[:pos], list_len * 4, fmt[pos + 1:])
  50.     # print(fmt, len(buf[offset:]), hex_str(buf[offset:]))
  51.     res = struct.unpack_from(fmt, buf, offset)
  52.     # return [x.decode() + '|' + hex_str(x) if isinstance(x, bytes) else x for x in res], fmt
  53.     return [fmt_str(x) if isinstance(x, bytes) else x for x in res], fmt
  54.     # return res, fmt
  55.  
  56.  
  57. class TTDSavegameDecoder:
  58.     def __init__(self):
  59.         pass
  60.  
  61.     @staticmethod
  62.     def _parse_label(data):
  63.         return ''.join(map(chr, data[:4]))
  64.  
  65.     @staticmethod
  66.     def read_gamma(data):
  67.         if (data[0] & 0xF0) == 0xE0:
  68.             return (data[0] & 0x0F) << 24 | data[1] << 16 | data[2] << 8 | data[3], 4
  69.         elif (data[0] & 0xE0) == 0xC0:
  70.             return (data[0] & 0x1F) << 16 | data[1] << 8 | data[2], 3
  71.         elif (data[0] & 0xC0) == 0x80:
  72.             return (data[0] & 0x3F) << 8 | data[1], 2
  73.         elif (data[0] & 0x80) == 0:
  74.             return data[0] & 0x7F, 1
  75.         raise ValueError('Invalid gamma')
  76.  
  77.     def _read_int(self, data, size, is_signed):
  78.         fmt = {8: '>b', 16: '>h', 32: '>i'}[size]
  79.         if not is_signed:
  80.             fmt = fmt.upper()
  81.         return struct.unpack_from(fmt, data)[0], size // 8
  82.  
  83.     def _decode_data(self, data, version):
  84.         d = memoryview(data)
  85.         # print('Savegame version', version, 'uncompressed size', len(d))
  86.         res = {}
  87.         while d:
  88.             chunk_id = self._parse_label(d[:4])
  89.             if chunk_id == '\0\0\0\0':
  90.                 break
  91.             assert chunk_id not in res, f"Duplicate chunk {chunk_id}"
  92.             # print('Chunk', chunk_id, end=' ')
  93.             block_mode = d[4]
  94.             d = d[5:]
  95.             # CH_RIFF         =  0,
  96.             # CH_ARRAY        =  1,
  97.             # CH_SPARSE_ARRAY =  2,
  98.             # CH_TYPE_MASK    =  3,
  99.             # CH_LAST         =  8, ///< Last chunk in this array.
  100.             # CH_AUTO_LENGTH  = 16,
  101.             if block_mode & 0xF == 0:  # RIFF
  102.                 size = d[0] << 16 | d[1] << 8 | d[2] | (block_mode >> 4) << 24;
  103.                 res[chunk_id] = d[3:3 + size]
  104.                 d = d[3 + size:]
  105.             elif block_mode == 1 or block_mode == 2:
  106.                 items = []
  107.                 indexes = []
  108.                 index = 0
  109.                 d0 = d
  110.                 dlen = 0
  111.                 while True:
  112.                     item_size, l = self.read_gamma(d)
  113.                     d = d[l:]
  114.                     dlen += l
  115.                     if not item_size:
  116.                         break
  117.                     item_size -= 1
  118.                     l = 0;
  119.                     if block_mode == 2:
  120.                         index, l = self.read_gamma(d)
  121.                     if item_size > 0:
  122.                         items.append(d[l:item_size])
  123.                     indexes.append(index)
  124.                     index += 1
  125.                     d = d[item_size:]
  126.                     dlen += item_size
  127.  
  128.                 res[chunk_id] = d0[:dlen]
  129.             else:
  130.                 raise ValueError(f'Unknown chunk mode {block_mode}')
  131.         return res
  132.  
  133.     def decode(self, data):
  134.         tag = self._parse_label(data)
  135.         decompressor = {
  136.             'OTTX': lzma.LZMADecompressor().decompress,
  137.             'OTTZ': zlib.decompress,
  138.         }.get(tag)
  139.         if not decompressor:
  140.             raise ValueError(f'Unknown savegame format {tag}')
  141.         major_version = data[4] * 8 + data[5]
  142.         minor_version = data[6] * 8 + data[7]
  143.         print(f'Version: {major_version} {minor_version}')
  144.  
  145.         res = {}
  146.         res['chunks'] = self._decode_data(decompressor(data[8:]), major_version)
  147.         res['version'] = major_version
  148.         return res
  149.  
  150.  
  151. d = TTDSavegameDecoder()
  152. s1 = d.decode(open(sys.argv[1], 'rb').read())
  153.  
  154. print(f'Version: {s1["version"]}')
  155.  
  156. for chunk_id, data in s1['chunks'].items():
  157.     print(chunk_id, hex_str(data[:100]))
  158.     l = len(data)
  159.     c = lzma.LZMACompressor(preset=2, check=lzma.CHECK_CRC32)
  160.     res = c.compress(data)
  161.     res += c.flush()
  162.     cl = len(res)
  163.     ratio = 100 * cl / l
  164.     print(f'{chunk_id} uncompressed: {l}  compressed: {cl}  ratio: {ratio:.1f}%')
  165.     print()
RAW Paste Data