RustNotMC/tools/extract-rustex-assets.py

#!/usr/bin/env python3
"""Extract the assets of the Rustex remake into a plain, decoded tree.

The remake ships its resources obfuscated: every file is XOR-ed with one
byte, and the byte differs per file.  The key is recoverable because the
first bytes of every format are known:

    PNG   89 50 4E 47      key = file[0] ^ 0x89
    OGG   4F 67 67 53      key = file[0] ^ 0x4F
    JSON  7B ...           key = file[0] ^ 0x7B

Files that are already plain (KTX2, MODL, the icon atlas inside game.jar)
are copied as they are.  The formats with a stronger cipher (.obj, .rrs,
.nbt, .properties, .lang) are copied untouched and reported, so that the
reverse engineering pass knows what is still missing.

The result lands in assets/rustex/ :

    legacy/    the Minecraft style assets of res/assets
    nextgen/   the high resolution models of res/nextgen
    gamejar/   the resources packed into game.jar (the icon atlas)
    maps/      the server map data (the monument layout)
    MANIFEST.json
"""

import json
import os
import shutil
import sys
import time
import zipfile
from pathlib import Path

SRC_ROOT = Path(r"G:\rustexremake\rustexremake")
SRC_RES = SRC_ROOT / "gamefiles" / "res"
SRC_JAR = SRC_ROOT / "gamefiles" / "game.jar"
SRC_MAPS = SRC_ROOT / "working_dir" / "rustex" / "maps"
DST = Path(r"G:\RustNotMC\assets\rustex")

PNG_MAGIC = b"\x89PNG\r\n\x1a\n"
OGG_MAGIC = b"OggS"
KTX2_MAGIC = b"\xabKTX 20\xbb\r\n\x1a\n"

# The first byte of the plain file, per extension.  None means "we do not
# know the format, try to brute force the key and score the result".
FIRST_BYTE = {
    ".png": 0x89,
    ".ogg": 0x4F,
    ".json": 0x7B,
}

# The formats we can never guess: their cipher is not a single byte.
OPAQUE = {".obj", ".rrs", ".nbt", ".bin", ".db", ".jpm"}

TEXTUAL = {".json", ".properties", ".lang", ".txt", ".mcmeta", ".fsh", ".vsh", ".glsl"}


def score_text(data: bytes) -> float:
    """How much of the data looks like readable text."""
    if not data:
        return 0.0
    good = 0
    for b in data:
        if b in (9, 10, 13) or 32 <= b < 127:
            good += 1
    return good / len(data)


def unxor(data: bytes, key: int) -> bytes:
    if key == 0:
        return data
    return bytes(b ^ key for b in data)


def decode_file(path: Path, ext: str):
    """Return (decoded bytes, key, kind, note)."""
    data = path.read_bytes()
    if not data:
        return data, 0, "empty", "пустой файл"

    # Already plain?
    if data.startswith(PNG_MAGIC):
        return data, 0, "png", ""
    if data.startswith(OGG_MAGIC):
        return data, 0, "ogg", ""
    if data.startswith(KTX2_MAGIC):
        return data, 0, "ktx2", ""
    if data.startswith(b"MODL"):
        return data, 0, "modl", ""
    if data.startswith(b"glTF"):
        return data, 0, "gltf", ""

    # The fast path: the key is implied by the first byte of the format.
    if ext in FIRST_BYTE:
        key = data[0] ^ FIRST_BYTE[ext]
        head = unxor(data[:16], key)
        if ext == ".png" and head.startswith(PNG_MAGIC):
            return unxor(data, key), key, "png", ""
        if ext == ".ogg" and head.startswith(OGG_MAGIC):
            return unxor(data, key), key, "ogg", ""
        if ext == ".json" and head[:1] in (b"{", b"["):
            return unxor(data, key), key, "json", ""

    # The brute force path: find the key that makes the file text.
    if ext in TEXTUAL or ext not in OPAQUE:
        best = (0.0, 0, b"")
        for key in range(1, 256):
            head = unxor(data[:512], key)
            s = score_text(head)
            if s > best[0]:
                best = (s, key, head)
        if best[0] >= 0.95:
            return unxor(data, best[1]), best[1], "text", ""

    # Nothing worked.
    return data, -1, "opaque", "ключ не подобрался"


def walk(base: Path):
    for dirpath, dirnames, filenames in os.walk(base):
        for name in filenames:
            yield Path(dirpath) / name


def extract_tree(src: Path, dst: Path, manifest: list, stats: dict):
    for path in walk(src):
        rel = path.relative_to(src)
        target = dst / rel
        ext = path.suffix.lower()
        try:
            decoded, key, kind, note = decode_file(path, ext)
        except Exception as exc:  # a broken file must not kill the run
            stats["error"] = stats.get("error", 0) + 1
            manifest.append({"src": str(rel), "status": "error", "note": str(exc)})
            continue

        target.parent.mkdir(parents=True, exist_ok=True)
        target.write_bytes(decoded)

        bucket = stats.setdefault(kind, {"count": 0, "bytes": 0, "encoded": 0})
        bucket["count"] += 1
        bucket["bytes"] += len(decoded)
        if key > 0:
            bucket["encoded"] += 1

        entry = {"src": str(rel), "kind": kind, "key": key, "bytes": len(decoded)}
        if note:
            entry["note"] = note
        manifest.append(entry)


def extract_jar(jar: Path, dst: Path, manifest: list, stats: dict):
    with zipfile.ZipFile(jar) as zf:
        for info in zf.infolist():
            if info.is_dir():
                continue
            name = info.filename
            if not name.startswith("assets/"):
                continue
            target = dst / name
            target.parent.mkdir(parents=True, exist_ok=True)
            data = zf.read(name)
            target.write_bytes(data)
            stats.setdefault("jar", {"count": 0, "bytes": 0, "encoded": 0})
            stats["jar"]["count"] += 1
            stats["jar"]["bytes"] += len(data)
            manifest.append({"src": name, "kind": "jar", "key": 0, "bytes": len(data)})


def main() -> int:
    started = time.time()
    if not SRC_RES.is_dir():
        print("нет исходника:", SRC_RES)
        return 1

    # The tree is overwritten file by file, so a leftover from an aborted run
    # is harmless.  A hard delete is only a convenience, and Windows likes to
    # keep a handle on a directory that was written a moment ago.
    shutil.rmtree(DST, ignore_errors=True)
    DST.mkdir(parents=True, exist_ok=True)

    manifest: list = []
    stats: dict = {}

    print("1/3 legacy assets ...", flush=True)
    extract_tree(SRC_RES / "assets", DST / "legacy", manifest, stats)

    print("2/3 nextgen models ...", flush=True)
    extract_tree(SRC_RES / "nextgen", DST / "nextgen", manifest, stats)

    print("3/3 game.jar + maps ...", flush=True)
    if SRC_JAR.is_file():
        extract_jar(SRC_JAR, DST / "gamejar", manifest, stats)
    if SRC_MAPS.is_dir():
        extract_tree(SRC_MAPS, DST / "maps", manifest, stats)

    total = sum(v["bytes"] for v in stats.values())
    print("")
    for kind in sorted(stats):
        v = stats[kind]
        print("  %-8s %6d файлов, %8.1f МБ, зашифровано %d"
              % (kind, v["count"], v["bytes"] / 1048576.0, v["encoded"]))
    print("всего: %.1f МБ за %.0f с" % (total / 1048576.0, time.time() - started))

    (DST / "MANIFEST.json").write_text(
        json.dumps({"stats": stats, "files": manifest}, ensure_ascii=False, indent=1),
        encoding="utf-8")
    print("манифест:", DST / "MANIFEST.json")
    return 0


if __name__ == "__main__":
    sys.exit(main())