Trainz 22 TZarc V2 decompressor

decompress.py

# trainz 22 tzarc v2 decompressor by Puyodead1
# as far as i know, this is only used in trainz plus but that might have changed? its used in archives from 2023 and later
# base information on the tzarc v1 format is from https://github.com/rileyzzz/ChumpLib/blob/main/chumplib/src/tzarcfile.cpp

import argparse
import zlib

# pip install binreader
from binreader import BinaryReader
# pip install tqdm
from tqdm import tqdm


def byte_swap(i: int) -> int:
    return ((i & 0x000000FF) << 24) | ((i & 0x0000FF00) << 8) | ((i & 0x00FF0000) >> 8) | ((i & 0xFF000000) >> 24)


def decompress(file_path: str):
    with open(file_path, "rb") as f:
        reader = BinaryReader(f)

        magic = reader.read_uint32()
        assert magic == 0x63615A54, "Invalid magic"

        version = reader.read_uint32()
        # byte swap
        version = byte_swap(version)
        assert version == 2, "Unsupported version"

        uncompressed_size = reader.read_uint32()
        unknown1 = reader.read_uint32()

        chunk_count = reader.read_uint32()
        chunk_sizes = [reader.read_uint16() + 1 for _ in range(chunk_count)]

        print(f"Uncompressed size: {uncompressed_size}")
        print(f"Unknown1: {unknown1}")
        print(f"Chunk count: {chunk_count}")
        # consumes too much time to print this, and who the fuck cares
        # for i, chunk_size in enumerate(chunk_sizes):
        #     print(f"\tChunk {i}: {chunk_size}")

        # read zlib chunks
        output_path = file_path.split(".")[0] + "_decompressed.tzarc"
        with open(output_path, "wb") as out:
            pbar = tqdm(total=chunk_count, unit="Chunks", unit_scale=True)
            # write a header for use in extraction
            out.write(b"TZac\x00\x00\x00\x01")

            decompressed_bytes = 0

            for i, chunk_size in enumerate(chunk_sizes):
                data = reader.read(chunk_size)
                decompressed = zlib.decompress(data)

                # print(f"Decompressed chunk {i+1}/{chunk_size}: {len(decompressed)} bytes")
                out.write(decompressed)
                out.flush()
                decompressed_bytes += len(decompressed)
                pbar.update(1)

            pbar.close()
            if decompressed_bytes == uncompressed_size:
                print(f"Decompression success! Wrote {decompressed_bytes} bytes to {file_path}.out")
            else:
                print(f"Decompression failed! Expected {uncompressed_size} bytes, got {decompressed_bytes} bytes")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Decompress TZArc V2 archives")
    parser.add_argument("file", help="The TZArc file to extract")

    args = parser.parse_args()
    decompress(args.file)