nnnotes 0.1.0__py3-none-any.whl

nnnotes/shader.py

@@ -0,0 +1,217 @@
+"""Shader extraction.
+
+Dumps every Shader object found in the given bundles:
+  - `<name>.json`                  parsed form: properties, subshaders, passes
+                                   (tags, render state), keyword names
+  - `<name>/<platform>/s<S>p<P>_<stage>_<N>.<ext>`
+                                   each compiled player sub-program (variant)
+  - `shaders.json`                 index with per-variant keywords
+
+Blob layout (Unity 2021.2+ / 6000.x): per platform, an LZ4 blob holding a table
+of (offset, length, segment) entries. Code entries and parameter-binding entries
+are separate; the parsed form's `m_PlayerSubPrograms` names, per variant, the
+code `m_BlobIndex`, its `m_KeywordIndices` and `m_GpuProgramType`. A code entry
+is: version, programType, 4 x int, keyword count + aligned strings, then the
+program bytes. GLES programs are GLSL ES text holding both stages under
+`#ifdef VERTEX` / `#ifdef FRAGMENT`. Vulkan programs are Unity's packed SPIR-V
+container and are written as-is (`.vkprog`).
+
+No shader is filtered by name; a viewer decides what it reimplements.
+
+A Shader object's dump depends on its serialized bytes only: it is made once per process (cache bucket
+"shaders", keyed by the object's bytes, type and Unity version) and written again from there.
+"""
+from __future__ import annotations
+
+import json
+import re
+import struct
+from pathlib import Path
+
+import UnityPy
+from UnityPy.enums import ShaderGpuProgramType
+from UnityPy.export.ShaderConverter import CompressionHelper
+
+from . import cache
+from .jsonio import dumps, write_json
+
+DUMPS = cache.bucket("shaders", salt=cache.source_salt(__file__, "UnityPy"), disk=True)
+
+# UnityEngine.Rendering ShaderCompilerPlatform value -> tag
+PLATFORM = {
+    0: "gl", 4: "d3d11", 5: "gles20", 9: "gles3", 14: "metal",
+    15: "glcore", 18: "vulkan", 19: "switch",
+}
+# which GPU program types belong to which platform
+PLATFORM_TYPES = {
+    9: {"GLES3", "GLES31", "GLES31AEP"},
+    5: {"GLES"},
+    15: {"GLCore32", "GLCore41", "GLCore43"},
+    18: {"SPIRV"},
+    14: {"MetalVS", "MetalFS"},
+}
+TEXT_TYPES = {"GLES", "GLES3", "GLES31", "GLES31AEP", "GLCore32", "GLCore41", "GLCore43",
+              "MetalVS", "MetalFS"}
+STAGES = ("progVertex", "progFragment", "progGeometry", "progHull", "progDomain")
+
+
+def _safe(name: str) -> str:
+    return re.sub(r"[^0-9A-Za-z._-]", "_", name)
+
+
+def _entry(arr, i):
+    x = arr[i]
+    return x[0] if isinstance(x, list) else x
+
+
+def _type_name(v: int) -> str:
+    return ShaderGpuProgramType(v).name.replace("kShaderGpuProgram", "")
+
+
+def _blob_table(raw: bytes) -> list[bytes]:
+    n = struct.unpack_from("<i", raw, 0)[0]
+    out = []
+    for i in range(n):
+        off, length, _seg = struct.unpack_from("<3i", raw, 4 + 12 * i)
+        out.append(raw[off:off + length])
+    return out
+
+
+def _code_of(entry: bytes) -> bytes:
+    """Program bytes of a code blob entry."""
+    pos = 4 + 4 + 16                           # version, programType, 4 x int
+    nkw = struct.unpack_from("<i", entry, pos)[0]
+    pos += 4
+    for _ in range(nkw):
+        n = struct.unpack_from("<i", entry, pos)[0]
+        pos += 4 + n
+        pos = (pos + 3) & ~3
+    n = struct.unpack_from("<i", entry, pos)[0]
+    return entry[pos + 4:pos + 4 + n]
+
+
+def platform_blobs(shader) -> dict[int, list[bytes]]:
+    blob = bytes(shader.compressedBlob or b"")
+    out = {}
+    n = min(len(shader.platforms), len(shader.compressedLengths), len(shader.decompressedLengths))
+    for i in range(n):
+        off = _entry(shader.offsets, i)
+        raw = CompressionHelper.decompress_lz4(
+            blob[off:off + _entry(shader.compressedLengths, i)],
+            _entry(shader.decompressedLengths, i))
+        out[shader.platforms[i]] = _blob_table(raw)
+    return out
+
+
+def variants(tt: dict, blobs: dict[int, list[bytes]]):
+    """Yield (platform, subshader, pass, stage, n, gpuType, keywords, code)."""
+    pf = tt["m_ParsedForm"]
+    names = pf.get("m_KeywordNames", [])
+    for si, ss in enumerate(pf["m_SubShaders"]):
+        for pi, ps in enumerate(ss["m_Passes"]):
+            for stage in STAGES:
+                prog = ps.get(stage) or {}
+                n = 0
+                for group in prog.get("m_PlayerSubPrograms", []):
+                    for sp in group:
+                        t = _type_name(sp["m_GpuProgramType"])
+                        for plat, table in blobs.items():
+                            if t not in PLATFORM_TYPES.get(plat, ()):
+                                continue
+                            code = _code_of(table[sp["m_BlobIndex"]])
+                            kws = [names[k] for k in sp.get("m_KeywordIndices", []) if k < len(names)]
+                            yield plat, si, pi, stage[4:].lower(), n, t, kws, code
+                        n += 1
+
+
+def _parsed_summary(tt: dict) -> dict:
+    pf = tt["m_ParsedForm"]
+    return {
+        "name": pf.get("m_Name"),
+        "properties": pf.get("m_PropInfo", {}).get("m_Props", []),
+        "keywords": pf.get("m_KeywordNames", []),
+        "subShaders": [{
+            "tags": ss.get("m_Tags"), "lod": ss.get("m_LOD"),
+            "passes": [{"name": ps.get("m_State", {}).get("m_Name") or ps.get("m_Name"),
+                        "tags": ps.get("m_Tags"), "state": ps.get("m_State")}
+                       for ps in ss["m_Passes"]],
+        } for ss in pf["m_SubShaders"]],
+        "fallback": pf.get("m_FallbackName"),
+    }
+
+
+def _render(o, tt: dict, name: str) -> tuple[str, list[dict], list[bytes]]:
+    """(parsed summary JSON text, variant records, variant program bytes) of one Shader object."""
+    safe = _safe(name)
+    text = dumps(_parsed_summary(tt), indent=1, ensure_ascii=False, default=str)
+    recs, codes = [], []
+    for plat, si, pi, stage, n, t, kws, code in variants(tt, platform_blobs(o.read())):
+        tag = PLATFORM.get(plat, str(plat))
+        ext = "glsl" if t in TEXT_TYPES and not t.startswith("Metal") else "metal" if t.startswith("Metal") else "vkprog"
+        fn = f"s{si}p{pi}_{stage}_{n}.{ext}"
+        recs.append({"file": f"{safe}/{tag}/{fn}", "platform": tag, "subShader": si, "pass": pi, "stage": stage,
+                     "type": t, "keywords": kws})
+        codes.append(bytes(code))
+    return text, recs, codes
+
+
+def _dump_key(o) -> str:
+    st = o.serialized_type
+    type_hash = bytes(st.old_type_hash) if st is not None and st.old_type_hash else b""
+    return DUMPS.key(o.get_raw_data(), type_hash, int(o.class_id), o.assets_file.unity_version,
+                     int(getattr(o, "platform", 0) or 0))
+
+
+def dump_objects(shaders, out_dir: Path, source: str, index: list | None = None) -> list:
+    """Write parsed summaries + every compiled variant of Shader objects.
+
+    `index` accumulates records across calls; a shader already present (same
+    name) is not written twice.
+    """
+    out_dir = Path(out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    index = [] if index is None else index
+    done = {r["name"] for r in index}
+    for o in shaders:
+        k = _dump_key(o)
+        hit = DUMPS.get(k)
+        if hit is not None:
+            meta, codes = cache.unpack(hit)
+            name, text, recs = json.loads(meta)
+        else:
+            tt = o.read_typetree()
+            named = tt.get("m_ParsedForm", {}).get("m_Name")
+            name = named or f"shader_{o.path_id}"
+            if name in done:
+                continue
+            text, recs, codes = _render(o, tt, name)
+            if named:                                  # a name from the path id is not a function of the bytes
+                DUMPS.put(k, cache.pack(json.dumps([name, text, recs], ensure_ascii=False).encode("utf-8"), codes))
+        if name in done:
+            continue
+        done.add(name)
+        safe = _safe(name)
+        (out_dir / f"{safe}.json").write_text(text, encoding="utf-8", newline="\n")      # as write_json writes it
+        rec = {"name": name, "source": source, "parsed": f"{safe}.json", "variants": []}
+        for v, code in zip(recs, codes):
+            pdir = out_dir / v["file"].rsplit("/", 1)[0]
+            pdir.mkdir(parents=True, exist_ok=True)
+            (out_dir / v["file"]).write_bytes(code)   # the program text or binary exactly as the game stores it
+            rec["variants"].append(v)
+        index.append(rec)
+    return index
+
+
+def write_index(index: list, out_dir: Path) -> dict:
+    write_json(Path(out_dir) / "shaders.json", index)
+    return {"count": len(index), "names": sorted(r["name"] for r in index),
+            "variants": sum(len(r["variants"]) for r in index)}
+
+
+def dump(bundles: list, out_dir: Path) -> dict:
+    index: list = []
+    for bpath in bundles:
+        env = UnityPy.load(str(bpath))
+        dump_objects([o for o in env.objects if o.type.name == "Shader"], out_dir,
+                     Path(bpath).name, index)
+    return write_index(index, out_dir)

nnnotes/spot.py

@@ -0,0 +1,201 @@
+"""Spot (explorable 2.5D scene) extractor -> spot.json + Spine assets.
+
+A spot = background prefab (cardboard room, see room.py) + situation prefab:
+SpotSituationSettings (camera/interaction parameters), SpotCharacterController,
+SpotCharacter x N (tap targets with focus transforms) and SpotSpineCharacter x M,
+each driving a Spine.Unity SkeletonAnimation. Components are identified by
+their MonoScript class (Catalog must be opened with `apk=`).
+
+A reference the prefab leaves empty is written as null, not filled in: a
+SpotSpineCharacter with no `_animation` has null `skeletonData`, `animation`
+and `world`; a SpotCharacter with no `_focus` has a null `focusWorld`.
+
+Spine files are written from the SkeletonDataAsset reference chain (write_skeleton):
+  SkeletonDataAsset.skeletonJSON           -> skeleton (.json or binary .skel)
+  SkeletonDataAsset.atlasAssets[i].atlasFile -> <name>.atlas
+  SpineAtlasAsset.materials[j]._MainTex    -> PNG named after atlas page j
+"""
+from __future__ import annotations
+
+import io
+import json
+from pathlib import Path
+
+import UnityPy
+
+from .catalog import Catalog
+from .config import apk_missing
+from .jsonio import write_json
+from .unity import SceneGraph, script_class, strip_pptrs, texture_image
+
+LANGS = ("_japanese", "_english", "_traditionalChinese", "_simplifiedChinese", "_korean")
+
+
+def _master(md: Path, name: str) -> list[dict]:
+    return json.loads((md / f"{name}.json").read_text(encoding="utf-8"))["_allData"]
+
+
+def _mat4(m) -> list[float]:
+    return [float(x) for x in m.T.reshape(-1)]          # column-major
+
+
+def _atlas_pages(text: str) -> list[str]:
+    """Page file names of a Spine atlas (a page line is followed by 'size:')."""
+    lines = [ln.strip() for ln in text.splitlines()]
+    return [ln for i, ln in enumerate(lines)
+            if ln and ":" not in ln and i + 1 < len(lines) and lines[i + 1].startswith("size:")]
+
+
+def _text_bytes(ta) -> bytes:
+    s = ta.m_Script
+    return s.encode("utf-8", "surrogateescape") if isinstance(s, str) else bytes(s)
+
+
+def _target(ref):
+    """The object a PPtr refers to, None for an empty reference (path id 0)."""
+    return ref.read() if ref.path_id else None
+
+
+def write_skeleton(o, write, written_atlases: dict[int, str] | None = None) -> dict:
+    """The Spine files of the SkeletonDataAsset object `o` (module docstring), each given to `write(file name,
+    bytes)` in this order: the skeleton, then for each of its atlases not in `written_atlases` ({SpineAtlasAsset
+    path id: atlas file name}, updated here) the atlas text as stored and one PNG per page. Returns the skeleton's
+    record of spot.json: name, skeleton file, atlas files, scale, default mix, mixes."""
+    written = {} if written_atlases is None else written_atlases
+    sda = o.read_typetree()
+    sk = o.read().skeletonJSON.read()
+    raw = _text_bytes(sk)
+    is_json = raw.lstrip()[:1] == b"{"
+    ext = ".json" if is_json else ".skel"
+    sk_file = sk.m_Name if sk.m_Name.endswith(ext) else f"{sk.m_Name}{ext}"   # as the atlas below
+    write(sk_file, raw)
+    atlases = []
+    for ap in o.read().atlasAssets:
+        atlas_asset = ap.read()
+        if ap.path_id not in written:
+            af = atlas_asset.atlasFile.read()
+            araw = _text_bytes(af)
+            atext = araw.decode("utf-8")
+            a_file = af.m_Name if af.m_Name.endswith(".atlas") else f"{af.m_Name}.atlas"
+            write(a_file, araw)                  # the atlas text exactly as the game stores it
+            pages = _atlas_pages(atext)
+            mats = list(atlas_asset.materials)
+            if len(mats) != len(pages):
+                raise RuntimeError(f"{a_file}: {len(pages)} pages vs {len(mats)} materials")
+            for page, mp in zip(pages, mats):
+                mat = mp.read()
+                tex = next(e.m_Texture.read() for k, e in mat.m_SavedProperties.m_TexEnvs
+                           if k == "_MainTex")
+                buf = io.BytesIO(); texture_image(tex).save(buf, format="PNG")
+                write(page, buf.getvalue())
+            written[ap.path_id] = a_file
+        atlases.append(written[ap.path_id])
+    return {
+        "name": sda["m_Name"], "skeleton": sk_file, "atlases": atlases,
+        "scale": sda["scale"], "defaultMix": sda["defaultMix"],
+        "mixes": [{"from": f, "to": t, "duration": d} for f, t, d in
+                  zip(sda.get("fromAnimation", []), sda.get("toAnimation", []), sda.get("duration", []))],
+    }
+
+
+def spine_characters(objs: list, graph, skeletons: dict, world_of) -> list[dict]:
+    """SpotSpineCharacter components: object path, the SkeletonAnimation they drive (its skeleton data name,
+    initial animation settings and world matrix). `skeletons`: SkeletonDataAsset path id -> skeleton record. One
+    whose `_animation` is empty has those three null."""
+    out = []
+    for o in objs:
+        ssc = o.read_typetree()
+        path = graph.path(graph.tf_of_go[ssc["m_GameObject"]["m_PathID"]])
+        anim_o = _target(o.read()._animation)
+        if anim_o is None:
+            out.append({"path": path, "skeletonData": None, "animation": None, "world": None})
+            continue
+        at = anim_o.object_reader.read_typetree()
+        out.append({
+            "path": path,
+            "skeletonData": skeletons[anim_o.skeletonDataAsset.path_id]["name"],
+            "animation": {k: at[k] for k in ("_animationName", "loop", "timeScale",
+                                             "initialSkinName", "initialFlipX", "initialFlipY",
+                                             "pmaVertexColors", "tintBlack", "zSpacing")},
+            "world": world_of(at["m_GameObject"]),
+        })
+    return out
+
+
+def tap_targets(objs: list, graph, world_of) -> list[dict]:
+    """SpotCharacter components (tap targets) by character id: object path, their fields, world matrix and the world
+    matrix of their `_focus` transform (null when the reference is empty)."""
+    out = []
+    for o in objs:
+        sc = o.read_typetree()
+        focus_o = _target(o.read()._focus)
+        out.append({
+            "path": graph.path(graph.tf_of_go[sc["m_GameObject"]["m_PathID"]]),
+            **strip_pptrs({k: v for k, v in sc.items() if k != "_focus"}),
+            "world": world_of(sc["m_GameObject"]),
+            "focusWorld": None if focus_o is None else world_of(focus_o.object_reader.read_typetree()["m_GameObject"]),
+        })
+    out.sort(key=lambda c: c["_characterId"])
+    return out
+
+
+def extract(cat: Catalog, md: Path, spot_id: int, out_dir: Path) -> dict:
+    if cat.apk is None:
+        raise apk_missing("resolving the Spot / Spine component classes")
+    spots = {r["_id"]: r for r in _master(md, "MasterHomeSpot")}
+    sp = spots[spot_id]
+    txt = {r["_id"]: r for r in _master(md, "MasterText")}
+    taps: dict[int, list[int]] = {}
+    for r in _master(md, "MasterStoryHomeSpotTapTalkEpisode"):
+        if r.get("_spotId") == spot_id:
+            taps.setdefault(r["_characterId"], []).append(r["_advId"])
+
+    out_dir = Path(out_dir)
+    spine_dir = out_dir / "spine"
+    spine_dir.mkdir(parents=True, exist_ok=True)
+
+    env = UnityPy.load(*[str(p) for p in cat.fetch_key(sp["_situationAssetPath"])])
+    graph = SceneGraph(env)
+    by_cls: dict[str, list] = {}
+    for o in env.objects:
+        if o.type.name == "MonoBehaviour":
+            by_cls.setdefault(script_class(o), []).append(o)
+
+    def world_of(go_pptr) -> list[float]:
+        return _mat4(graph.world_of_go(go_pptr["m_PathID"]))
+
+    settings = strip_pptrs(by_cls["SpotSituationSettings"][0].read_typetree())
+    controller_o = by_cls["SpotCharacterController"][0]
+    controller = strip_pptrs(controller_o.read_typetree())
+
+    # Spine skeleton data -> files
+    def write(name: str, data: bytes) -> None:
+        (spine_dir / name).write_bytes(data)
+
+    skeleton_files: dict[int, dict] = {}
+    written_atlases: dict[int, str] = {}
+    for o in by_cls.get("SkeletonDataAsset", []):
+        skeleton_files[o.path_id] = write_skeleton(o, write, written_atlases)
+
+    spine_chars = spine_characters(by_cls.get("SpotSpineCharacter", []), graph, skeleton_files, world_of)
+    characters = tap_targets(by_cls.get("SpotCharacter", []), graph, world_of)
+
+    doc = {
+        "spotId": spot_id,
+        "master": sp,
+        "name": {lang[1:]: txt.get(sp.get("_nameTextId"), {}).get(lang, "") for lang in LANGS},
+        "tapTalk": {str(k): v for k, v in sorted(taps.items())},
+        "situationSettings": settings,
+        "controller": {k: v for k, v in controller.items()
+                       if k not in ("_spotSpineCharacters", "_characters")},
+        "characters": characters,
+        "spineCharacters": spine_chars,
+        "skeletons": list(skeleton_files.values()),
+        "resources": {
+            "background": [b.name for b in cat.resolve(sp["_backgroundAssetPath"])],
+            "situation": [b.name for b in cat.resolve(sp["_situationAssetPath"])],
+        },
+        "matrixConvention": "Unity world space (left-handed), 4x4 column-major",
+    }
+    write_json(out_dir / "spot.json", doc)
+    return doc