local-translator/local_translator/unity_scripts.py

"""Strict codec for Unity 2019 managed-reference scenario objects.

Schemas come from the game's own assemblies. The old version-1 registry uses
implicit integer IDs and a Terminus record, rather than the modern array format.
Use in an isolated asset-processing process; UnityPy's codec hooks are global.
"""
from contextlib import contextmanager


class ScriptCodec:
    def __init__(self, game, version="2019.4.22f1"):
        from UnityPy.helpers.TypeTreeGenerator import TypeTreeGenerator
        self.generator = TypeTreeGenerator(version)
        self.generator.load_local_game(str(game))
        self.cache = {}

    def node(self, assembly, fullname, reference=False):
        from UnityPy.helpers import TypeTreeHelper as helper
        from UnityPy.helpers.TypeTreeNode import TypeTreeNode
        key = (assembly, fullname, reference)
        if key in self.cache:
            return self.cache[key]
        source = self.generator.get_nodes_up(assembly, fullname)
        # The generator calls List<string> "string"; the generic reader would
        # mistake its count for a scalar string length. Normalize array nodes.
        def clone(n):
            values = n.to_dict()
            values["m_Children"] = [clone(c) for c in n.m_Children]
            return TypeTreeNode(**values)
        node = clone(source)
        for child in node.traverse():
            if (child.m_Children and child.m_Children[0].m_Type == "Array"
                    and child.m_Type in helper.FUNCTION_READ_MAP
                    and child.m_Children[0].m_Children[1].m_Type != "char"):
                child.m_Type = "vector"
        if reference:
            # Generator assumes MonoBehaviour even for plain serializable classes.
            children = [c for c in node.m_Children if c.m_Name not in
                        {"m_GameObject", "m_Enabled", "m_Script", "m_Name", "references"}]
            node = TypeTreeNode(0, node.m_Type, node.m_Name, -1, 0, m_Children=children)
        self.cache[key] = node
        return node

    def ref_node(self, typ):
        if not typ["class"]:
            return None
        if not typ["asm"] or not all(ord(c) >= 32 for c in typ["class"] + typ["ns"] + typ["asm"]):
            raise ValueError("Invalid managed-reference type; refusing to guess its layout")
        return self.node(typ["asm"], ".".join(filter(None, [typ["ns"], typ["class"]])), True)

    @contextmanager
    def active(self):
        from UnityPy.helpers import TypeTreeHelper as helper
        old_read, old_write = helper.read_value, helper.write_value
        old_boost, old_resolve = helper.read_typetree_boost, helper.get_ref_type_node

        def read(node, reader, config):
            if node.m_Type != "ManagedReferencesRegistry":
                return old_read(node, reader, config)
            version = reader.read_int()
            if version != 1:
                raise ValueError(f"Unsupported managed-reference registry: {version}")
            type_node = node.m_Children[1].m_Children[0]
            records = []
            while len(records) < 1_000_000:
                typ = old_read(type_node, reader, config)
                if typ == {"class": "Terminus", "ns": "UnityEngine.DMAT", "asm": "FAKE_ASM"}:
                    return {"version": version, "records": records, "terminus": typ}
                schema = self.ref_node(typ)
                records.append({"type": typ, "data": read(schema, reader, config) if schema else {}})
            raise ValueError("Unterminated managed-reference registry")

        def write(value, node, writer, config):
            if node.m_Type != "ManagedReferencesRegistry":
                return old_write(value, node, writer, config)
            if value["version"] != 1:
                raise ValueError("Unsupported registry version")
            writer.write_int(1)
            type_node = node.m_Children[1].m_Children[0]
            for ref in value["records"]:
                old_write(ref["type"], type_node, writer, config)
                schema = self.ref_node(ref["type"])
                if schema:
                    write(ref["data"], schema, writer, config)
            old_write(value["terminus"], type_node, writer, config)

        helper.read_typetree_boost = None
        helper.get_ref_type_node = lambda ref, _: self.ref_node(ref["type"])
        helper.read_value, helper.write_value = read, write
        try:
            yield self
        finally:
            helper.read_value, helper.write_value = old_read, old_write
            helper.read_typetree_boost, helper.get_ref_type_node = old_boost, old_resolve

    def read(self, obj):
        from UnityPy.streams import EndianBinaryWriter
        from UnityPy.helpers import TypeTreeHelper as helper
        node = self.node("Elringus.Naninovel.Runtime", "Naninovel.Script")
        tree = obj.read_typetree(nodes=node)
        writer = EndianBinaryWriter(endian=obj.reader.endian)
        helper.write_typetree(tree, node, writer, obj.assets_file)
        if writer.bytes != obj.get_raw_data():
            raise ValueError(f"Script {obj.path_id}: binary round-trip mismatch")
        return tree


def script_objects(environment):
    for obj in environment.objects:
        if obj.type.name != "MonoBehaviour":
            continue
        try:
            head = obj.parse_monobehaviour_head()
            script = head.m_Script.read()
        except (ValueError, AttributeError, FileNotFoundError):
            continue
        if script.m_Namespace == "Naninovel" and script.m_ClassName == "Script":
            yield obj