revelation.py: extract .agents files

revelation.py

import errno
import os
import re
import struct
import sys
import zlib

if sys.version_info.major == 2:
    from StringIO import StringIO as BytesIO
else:
    from io import BytesIO

DEBUG = False

TAG_BLOCKS = {
    "AGNT",  # C3 agent
    "DSAG",  # DS agent
    "MACH",  # SM agent
    "HAND",  # SM agent
    "LIVE",  # SM agent
    "MONK",  # SM agent
    "EXPC",  # C3 creature info
    "DSEX",  # DS creature info
    "SFAM",  # C3 starter family
    "EGGS",  # eggs
    "DFAM",  # DS starter family
}


def write_binary(data, filename):
    print("Writing {}".format(filename))
    with open(filename, "wb") as f:
        f.write(data)


def write_string(s, filename):
    print("Writing {}".format(filename))
    with open(filename, "wb") as f:
        f.write(s.encode("cp1252"))


def read_or_fail(f, num_bytes):
    data = f.read(num_bytes)
    if len(data) != num_bytes:
        raise Exception(
            "Expected to read {} bytes, but only got {} bytes".format(
                num_bytes, len(data)
            )
        )
    return data


def readu32le(f):
    return struct.unpack("<I", read_or_fail(f, 4))[0]


def escape(s):
    ret = ""
    for c in s:
        if c == "\n":
            ret += "\\n"
        elif c == "\r":
            ret += "\\r"
        elif c == "\t":
            ret += "\\t"
        elif c == '"':
            ret += '\\"'
        else:
            ret += c
    return ret


def main():
    if len(sys.argv) != 2:
        sys.stderr.write("USAGE: {} FILE\n".format(sys.argv[0]))
        exit(1)
    filename = sys.argv[1]
    filenamestem = os.path.splitext(os.path.basename(filename))[0]

    with open(filename, "rb") as f:
        magic = read_or_fail(f, 4)

        if magic != b"PRAY":
            raise Exception(
                "Not a PRAY file: expected file magic 'PRAY', got {}".format(
                    repr(magic)
                )
            )

        output_directory = filenamestem
        try:
            os.makedirs(output_directory)
        except OSError as e:
            if e.errno != errno.EEXIST:
                raise
        pray_text = '"en-GB"\n\n'

        big_string_values_to_filenames = {}
        inline_files = []

        while True:
            block_type = f.read(4).decode("cp1252")
            if len(block_type) == 0:
                break
            if len(block_type) != 4:
                raise Exception(
                    "Expected to read 4 bytes, but only got {} bytes".format(
                        len(block_type)
                    )
                )

            block_name = read_or_fail(f, 128).split(b"\x00")[0].decode("cp1252")
            length_compressed = readu32le(f)
            length_decompressed = readu32le(f)
            flags = readu32le(f)
            if DEBUG:
                print(
                    "block type={} name={} length_compressed={} length_decompressed={} flags={}".format(
                        repr(block_type),
                        repr(block_name),
                        length_compressed,
                        length_decompressed,
                        flags,
                    )
                )
            if flags not in (0, 1):
                raise Exception(
                    "Expected block flags to be 0 (uncompressed) or 1 (compressed), got {}".format(
                        hex(flags)
                    )
                )
            is_compressed = flags == 1

            if not is_compressed and length_compressed != length_decompressed:
                sys.stderr.write(
                    "WARNING: compression flag not set but lengths don't match, assuming data is actually compressed\n"
                )
                is_compressed = True

            data = read_or_fail(f, length_compressed)
            if is_compressed:
                data = zlib.decompress(data)
                if len(data) != length_decompressed:
                    sys.stderr.write(
                        "WARNING: data size after decompression doesn't match size specified in header, expected {} but got {}\n".format(
                            length_decompressed, len(data)
                        )
                    )

            if block_type in TAG_BLOCKS:
                reader = BytesIO(data)

                int_tags = []
                num_int_tags = readu32le(reader)
                if DEBUG:
                    print("num_int_tags = {}".format(num_int_tags))
                for _ in range(num_int_tags):
                    name_length = readu32le(reader)
                    name = read_or_fail(reader, name_length).decode("cp1252")
                    value = readu32le(reader)
                    int_tags += [(name, value)]
                    if DEBUG:
                        print((name, value))

                string_tags = []
                num_string_tags = readu32le(reader)
                if DEBUG:
                    print("num_string_tags = {}".format(num_string_tags))
                for _ in range(num_string_tags):
                    name_length = readu32le(reader)
                    name = read_or_fail(reader, name_length).decode("cp1252")
                    value_length = readu32le(reader)
                    value = read_or_fail(reader, value_length).decode("cp1252")
                    string_tags += [(name, value)]
                    if DEBUG:
                        print((name, value))

                pray_text += 'group {} "{}"\n'.format(block_type, block_name)
                for (k, v) in int_tags:
                    pray_text += '"{}" {}\n'.format(k, v)
                for (k, v) in string_tags:
                    if re.search(r"^(Script \d+)|(Remove script)$", k) and "\n" in v:
                        if v not in big_string_values_to_filenames:
                            tag_filename = block_name + " - " + k + ".cos"
                            write_string(
                                v, os.path.join(output_directory, tag_filename)
                            )
                            big_string_values_to_filenames[v] = tag_filename
                        pray_text += '"{}" @ "{}"\n'.format(
                            k, big_string_values_to_filenames[v]
                        )
                    else:
                        pray_text += '"{}" "{}"\n'.format(k, escape(v))
                pray_text += "\n"
            else:
                write_binary(data, os.path.join(output_directory, block_name))
                inline_files.append((block_type, block_name))

    for block_type, block_name in sorted(inline_files):
        pray_text += 'inline {} "{}" "{}"\n'.format(block_type, block_name, block_name)

    write_string(pray_text, os.path.join(output_directory, filenamestem + ".ps.txt"))


if __name__ == "__main__":
    main()