Factorio map metadata parser

save.py

from zipfile import ZipFile
from struct import Struct


## primitives ##

class Deserializer:
    u16 = Struct('<H')
    u32 = Struct('<I')

    def __init__(self, stream):
        self.stream = stream
        self.version = tuple(self.read_u16() for _ in range(4))

    def read(self, n):
        res = self.stream.read(n)
        assert len(res) == n
        return res

    def read_fmt(self, fmt):
        return fmt.unpack(self.read(fmt.size))[0]

    def read_optim(self, dtype):
        if self.version >= (0, 14, 14, 0):
            byte = self.read_u8()
            if byte != 0xFF:
                return byte
        return self.read_fmt(dtype)

    # public API

    def read_u8(self):
        return self.read(1)[0]

    def read_bool(self):
        return { 0: False, 1: True }[self.read_u8()]

    def read_u16(self):
        return self.read_fmt(self.u16)

    def read_u32(self):
        return self.read_fmt(self.u32)

    def read_optim_u16(self):
        return self.read_optim(self.u16)

    def read_optim_u32(self):
        return self.read_optim(self.u32)

    def read_length(self):
        if self.version >= (0, 16, 0, 0):
            return self.read_optim_u32()
        return self.read_u32()

    def read_str(self):
        length = self.read_length()
        return self.read(length).decode('utf-8')

    def read_optim_str(self):
        length = self.read_optim_u32()
        return self.read(length).decode('utf-8')

    def read_array_fixed(self, func, length):
        return [ func(self) for _ in range(length) ]

    def read_array(self, func, length_func=None):
        return self.read_fixed( func, (length_func or self.read_u32)() )


## level.dat structures ##

def read_version(ds: Deserializer):
    return tuple( ds.read_array_fixed(lambda ds: ds.read_optim_u16(), 3) )

def read_stats(ds: Deserializer):
    force_id = ds.read_u8()
    stats = ds.read_array_fixed(read_stats_item, 3)
    return force_id, stats

def read_stats_item(ds: Deserializer):
    return ds.read_array(lambda ds: (ds.read_u16(), ds.read_u32()) )

def read_mod(ds: Deserializer):
    name = ds.read_optim_str()
    value = {}
    value['version'] = read_version(ds)
    if ds.version > (0, 15, 0, 91):
        value['crc'] = ds.read_u32()
    return name, value



class SaveFile:
    def __init__(self, filename):
        zf = ZipFile(filename, 'r')
        datfile = None

        for f in zf.namelist():
            if f.endswith('/level.dat'):
                datfile = f
                break

        if not datfile:
            raise IOError("level.dat not found in save file")

        self.read_level( Deserializer(zf.open(datfile)) )

    def read_level(self, ds: Deserializer):
        self.campaign = ds.read_str()
        self.name = ds.read_str()
        self.base_mod = ds.read_str()

        if ds.version > (0, 17, 0, 0):
            self.base = ds.read_str()

        self.difficulty = ds.read_u8()  # 0: Normal, 1: Old School, 2: Hardcore

        self.finished = ds.read_bool()
        self.player_won = ds.read_bool()
        self.next_level = ds.read_str()  # usually empty

        if ds.version >= (0, 12, 0, 0):
            self.can_continue = ds.read_bool()
            self.finished_but_continuing = ds.read_bool()

        self.saving_replay = ds.read_bool()

        if ds.version >= (0, 16, 0, 0):
            self.allow_non_admin_debug_options = ds.read_bool()

        self.loaded_from = read_version(ds)
        self.loaded_from_build = ds.read_u16()

        self.allowed_commands = ds.read_u8()
        self.allowed_commands_normalized = self.allowed_commands
        if ds.version <= (0, 13, 0, 87):
            self.allowed_commands_normalized = 1 if self.allowed_commands else 2

        if ds.version <= (0, 13, 0, 42):
            self.stats = ds.read_array(read_stats)

        self.mods = ds.read_array(read_mod, ds.read_length)


if __name__ == '__main__':
    import sys
    try:
        from yaml import safe_dump
    except ImportError:
        print('Install PyYAML for pretty printing')

        def safe_dump(s, **kw):
            return repr(s)

    for name in sys.argv[1:]:
        sf = SaveFile(name)
        print('%s:' % name)
        print()
        print(safe_dump(sf.__dict__, default_flow_style=False))
        print('---')