pokemon_pinball_wiki/scripts/extract_godot_pck_assets.py

#!/usr/bin/env python3
"""Extract and convert assets from the Pokemon Pro Godot PCK.

This is a static extractor. It does not run the target AArch64 game binary.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import math
import os
import re
import shutil
import struct
import subprocess
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Optional


PCK_TABLE_OFFSET = 0x64
PCK_FILE_BASE_OFFSET = 0x18
PCK_FILE_COUNT_OFFSET = 0x60
HEX32_RE = re.compile(r"^(?P<name>.+)-(?P<hash>[0-9a-fA-F]{32})\.ctex$")


@dataclass
class PckEntry:
    path: str
    offset: int
    size: int
    md5: str
    flags: int


def sha256_file(path: Path) -> str:
    h = hashlib.sha256()
    with path.open("rb") as f:
        for chunk in iter(lambda: f.read(1024 * 1024), b""):
            h.update(chunk)
    return h.hexdigest()


def safe_parts(path: str) -> list[str]:
    parts: list[str] = []
    for part in Path(path).parts:
        if part in ("", ".", ".."):
            continue
        clean = re.sub(r"[^A-Za-z0-9._' -]+", "_", part)
        clean = clean.strip()
        parts.append(clean or "_")
    return parts


def safe_join(root: Path, path: str) -> Path:
    out = root.joinpath(*safe_parts(path))
    resolved = out.resolve()
    root_resolved = root.resolve()
    if root_resolved != resolved and root_resolved not in resolved.parents:
        raise ValueError(f"unsafe output path: {path}")
    return out


def parse_pck(data: bytes) -> tuple[int, list[PckEntry]]:
    if data[:4] != b"GDPC":
        raise ValueError("not a Godot PCK: missing GDPC magic")
    file_base = struct.unpack_from("<Q", data, PCK_FILE_BASE_OFFSET)[0]
    file_count = struct.unpack_from("<I", data, PCK_FILE_COUNT_OFFSET)[0]
    pos = PCK_TABLE_OFFSET
    entries: list[PckEntry] = []
    for _ in range(file_count):
        path_len = struct.unpack_from("<I", data, pos)[0]
        pos += 4
        raw_path = data[pos : pos + path_len]
        pos += path_len
        path = raw_path.rstrip(b"\0").decode("utf-8", "replace")
        offset, size = struct.unpack_from("<QQ", data, pos)
        md5 = data[pos + 16 : pos + 32].hex()
        flags = struct.unpack_from("<I", data, pos + 32)[0]
        pos += 36
        entries.append(PckEntry(path=path, offset=offset, size=size, md5=md5, flags=flags))
    if pos != file_base:
        raise ValueError(f"PCK table ended at 0x{pos:x}, file_base is 0x{file_base:x}")
    return file_base, entries


def decompress_rscc(blob: bytes, zstd_path: str) -> bytes:
    if not blob.startswith(b"RSCC"):
        return blob
    if len(blob) < 16:
        raise ValueError("truncated RSCC header")
    _version, block_size, total_size = struct.unpack_from("<III", blob, 4)
    if block_size == 0:
        raise ValueError("invalid RSCC block size 0")
    block_count = math.ceil(total_size / block_size)
    size_table_end = 16 + (4 * block_count)
    if len(blob) < size_table_end:
        raise ValueError("truncated RSCC block table")
    block_sizes = struct.unpack_from("<" + "I" * block_count, blob, 16)
    compressed_start = size_table_end
    compressed_end = compressed_start + sum(block_sizes)
    compressed = blob[compressed_start:compressed_end]
    result = subprocess.run(
        [zstd_path, "-dc"],
        input=compressed,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        check=False,
    )
    if result.returncode != 0:
        err = result.stderr.decode("utf-8", "replace").strip()
        raise ValueError(f"zstd failed: {err}")
    if len(result.stdout) != total_size:
        raise ValueError(f"RSCC decompressed to {len(result.stdout)} bytes, expected {total_size}")
    return result.stdout


def extract_riff_webp(data: bytes) -> Optional[bytes]:
    idx = 0
    while True:
        idx = data.find(b"RIFF", idx)
        if idx < 0:
            return None
        if idx + 12 <= len(data) and data[idx + 8 : idx + 12] == b"WEBP":
            riff_size = struct.unpack_from("<I", data, idx + 4)[0] + 8
            end = idx + riff_size
            if end <= len(data):
                return data[idx:end]
        idx += 4


def extract_jpeg(data: bytes) -> Optional[bytes]:
    start = data.find(b"\xff\xd8\xff")
    if start < 0:
        return None
    end = data.find(b"\xff\xd9", start + 2)
    if end < 0:
        return None
    return data[start : end + 2]


def extract_ttf(data: bytes) -> Optional[bytes]:
    candidates = [data.find(b"\x00\x01\x00\x00"), data.find(b"OTTO"), data.find(b"ttcf")]
    candidates = [idx for idx in candidates if idx >= 0]
    if not candidates:
        return None
    start = min(candidates)
    if data[start : start + 4] == b"ttcf":
        return data[start:]
    if start + 12 > len(data):
        return None
    num_tables = struct.unpack_from(">H", data, start + 4)[0]
    table_dir_end = start + 12 + (num_tables * 16)
    if num_tables == 0 or table_dir_end > len(data):
        return None
    max_end = table_dir_end
    for i in range(num_tables):
        record = start + 12 + (i * 16)
        table_offset = struct.unpack_from(">I", data, record + 8)[0]
        table_length = struct.unpack_from(">I", data, record + 12)[0]
        max_end = max(max_end, start + table_offset + table_length)
    if max_end > len(data):
        return None
    return data[start:max_end]


def image_output_name(entry_path: str, extension: str) -> str:
    name = Path(entry_path).name
    match = HEX32_RE.match(name)
    if match:
        source_name = match.group("name")
        short_hash = match.group("hash")[:8].lower()
        stem = Path(source_name).stem
        return f"{stem}--{short_hash}{extension}"
    return f"{Path(name).stem}{extension}"


def text_like(data: bytes) -> bool:
    if b"\0" in data[:4096]:
        return False
    try:
        data.decode("utf-8")
        return True
    except UnicodeDecodeError:
        return False


def write_file(path: Path, data: bytes) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_bytes(data)


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--pck", type=Path, default=Path("games/pokemon_pro/assets/godot/main.pck"))
    parser.add_argument("--output", type=Path, default=Path("assets"))
    args = parser.parse_args()

    zstd_path = shutil.which("zstd") or "/opt/homebrew/bin/zstd"
    if not Path(zstd_path).exists():
        print("zstd is required to decompress RSCC resources", file=sys.stderr)
        return 1
    if args.output.exists():
        print(f"output path already exists: {args.output}", file=sys.stderr)
        return 1

    data = args.pck.read_bytes()
    file_base, entries = parse_pck(data)

    raw_root = args.output / "godot_raw"
    openable_root = args.output / "openable"
    stats = {
        "entries": len(entries),
        "raw_files": 0,
        "rscc_decompressed": 0,
        "images_webp": 0,
        "images_jpeg": 0,
        "videos_mp4": 0,
        "fonts_ttf": 0,
        "text_files": 0,
        "conversion_failures": [],
    }
    manifest_entries = []

    for entry in entries:
        start = file_base + entry.offset
        end = start + entry.size
        blob = data[start:end]
        try:
            decoded = decompress_rscc(blob, zstd_path)
            if decoded is not blob:
                stats["rscc_decompressed"] += 1
        except Exception as exc:
            decoded = blob
            stats["conversion_failures"].append({"path": entry.path, "stage": "rscc", "error": str(exc)})

        raw_path = safe_join(raw_root, entry.path)
        write_file(raw_path, decoded)
        stats["raw_files"] += 1

        lower_path = entry.path.lower()
        converted_path: Optional[Path] = None

        if lower_path.endswith(".ctex"):
            webp = extract_riff_webp(decoded)
            if webp is not None:
                converted_path = openable_root / "images" / "imported" / image_output_name(entry.path, ".webp")
                write_file(converted_path, webp)
                stats["images_webp"] += 1
            else:
                jpeg = extract_jpeg(decoded)
                if jpeg is not None:
                    converted_path = openable_root / "images" / "imported" / image_output_name(entry.path, ".jpg")
                    write_file(converted_path, jpeg)
                    stats["images_jpeg"] += 1
                else:
                    stats["conversion_failures"].append({"path": entry.path, "stage": "image-carve", "error": "no WEBP/JPEG payload found"})
        elif lower_path.endswith(".fontdata"):
            ttf = extract_ttf(decoded)
            if ttf is not None:
                name = Path(entry.path).name.split("-", 1)[0]
                converted_path = openable_root / "fonts" / name
                write_file(converted_path, ttf)
                stats["fonts_ttf"] += 1
            else:
                stats["conversion_failures"].append({"path": entry.path, "stage": "font-carve", "error": "no sfnt font payload found"})
        elif lower_path.endswith(".mp4"):
            converted_path = safe_join(openable_root / "videos", entry.path)
            write_file(converted_path, decoded)
            stats["videos_mp4"] += 1
        elif lower_path.endswith((".txt", ".cfg", ".import", ".remap")) or text_like(decoded):
            converted_path = safe_join(openable_root / "text", f"{entry.path}.txt")
            write_file(converted_path, decoded)
            stats["text_files"] += 1

        manifest_entries.append(
            {
                "path": entry.path,
                "offset": entry.offset,
                "size": entry.size,
                "md5": entry.md5,
                "flags": entry.flags,
                "decoded_size": len(decoded),
                "raw_path": str(raw_path.relative_to(args.output)),
                "openable_path": str(converted_path.relative_to(args.output)) if converted_path else None,
            }
        )

    manifest = {
        "source": str(args.pck),
        "source_sha256": sha256_file(args.pck),
        "source_size": args.pck.stat().st_size,
        "pck_file_base": file_base,
        "pck_entry_count": len(entries),
        "stats": stats,
        "entries": manifest_entries,
    }
    args.output.mkdir(parents=True, exist_ok=True)
    (args.output / "manifest.json").write_text(json.dumps(manifest, indent=2), encoding="utf-8")
    (args.output / "README.md").write_text(
        "# Extracted Pokemon Pro Godot Assets\n\n"
        "Generated from `/games/pokemon_pro/assets/godot/main.pck` by "
        "`rootfs-triage-wiki/scripts/extract_godot_pck_assets.py`.\n\n"
        "- `godot_raw/` preserves every PCK entry after RSCC/zstd decompression where applicable.\n"
        "- `openable/images/imported/` contains WebP/JPEG payloads carved from Godot `.ctex` textures.\n"
        "- `openable/videos/` contains MP4 assets from the pack.\n"
        "- `openable/fonts/` contains TTF payloads carved from Godot `.fontdata` resources.\n"
        "- `openable/text/` contains text metadata and config files with `.txt` suffixes for easy opening.\n"
        "- `manifest.json` records offsets, sizes, hashes, decoded paths, and conversion results.\n",
        encoding="utf-8",
    )
    print(json.dumps({"output": str(args.output), "stats": stats}, indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())