Kilohearts Content Bank File extractor

extract.py

# Kilohearts Content Bank Extractor
# Puyodead1
#
# Why? idk, because I can

import argparse
import os
import pathlib

from binreader import BinaryReader
from construct import *

MAGIC = b"\x89kHs\r\n\x1a\nBank0001"

class IndexEntry(object):
    def __init__(self, unknown, offset, size):
        self.unknown = unknown
        self.offset = offset
        self.size = size

    @staticmethod
    def parse(reader):
        unknown = reader.read_int64()
        offset = reader.read_int64()
        size = reader.read_int64()
        return IndexEntry(unknown, offset, size)


def extract(file: str):
    if not os.path.isfile(file):
        raise ValueError("File does not exist")

    a = pathlib.Path(file).name.split(".")

    if len(a) != 3 or a[0] != "kilohearts" or a[2] != "bank":
        raise ValueError("Invalid kilohearts content bank file")

    with open(file, "rb") as f:
        reader = BinaryReader(f)

        magic = reader.read(16)
        if magic != MAGIC:
            raise ValueError("Invalid kilohearts content bank file, unexpected magic")

        entry_count = reader.read_int64()
        data_directory = [IndexEntry.parse(reader) for _ in range(entry_count)]
        string_count = reader.read_int64()
        string_directory = [reader.read_cstring() for _ in range(string_count)]

        root_dir = pathlib.Path("output", a[1])
        root_dir.mkdir(parents=True, exist_ok=True)

        for i in range(entry_count):
            dd = data_directory[i]
            se = string_directory[i].decode()
            print(f"Entry {i}: {se} @ {dd.offset} ({dd.size} bytes) (unknown: {dd.unknown})")

            # if the offset is 0, it's a directory (and size is 0 too)
            if dd.offset == 0:
                dir = pathlib.Path(root_dir, se)
                dir.mkdir(parents=True, exist_ok=True)
            else:
                with open(pathlib.Path(root_dir, se), "wb") as f:
                    # seek to the offset
                    reader.seek(dd.offset)
                    # read the data
                    f.write(reader.read(dd.size))


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Extracts the contents of a kilohearts content bank file")
    parser.add_argument("file", help="Path to bank file")
    args = parser.parse_args()
    extract(args.file)