MonoBehaviour nodes from TypeTreeGenerator misread m_Script and change it on save #388

import importlib.metadata
import struct
import sys

import UnityPy
from UnityPy.enums import ClassIDType
from UnityPy.helpers.Tpk import get_typetree_node
from UnityPy.helpers.TypeTreeGenerator import TypeTreeGenerator
from UnityPy.helpers.TypeTreeHelper import read_typetree, write_typetree
from UnityPy.helpers.TypeTreeNode import TypeTreeNode
from UnityPy.helpers.UnityVersion import UnityVersion
from UnityPy.streams import EndianBinaryReader, EndianBinaryWriter

# Game.dll holds:  public class Sample : MonoBehaviour { public int count; public bool flag; public int after; }
VERSION = "2022.3.0f1"
print("UnityPy", UnityPy.__version__, "TypeTreeGeneratorAPI", importlib.metadata.version("TypeTreeGeneratorAPI"))


def header(node):
    flags = []
    for child in node.m_Children:
        flags.append(f"{child.m_Name}{' (align)' if child.m_MetaFlag & 0x4000 else ''}")
        if child.m_Name == "m_Name":
            return ", ".join(flags)


def aligned_string(value):
    data = value.encode()
    return struct.pack("<i", len(data)) + data + bytes(-len(data) % 4)


# a Sample MonoBehaviour as Unity serializes it: m_Enabled is followed by 3 bytes of padding
data = (
    struct.pack("<iq", 0, 1)  # m_GameObject
    + struct.pack("<B3x", 1)  # m_Enabled
    + struct.pack("<iq", 0, 0x1122334455667788)  # m_Script
    + aligned_string("sample")  # m_Name
    + struct.pack("<i", 5)  # count
    + struct.pack("<B3x", 1)  # flag
    + struct.pack("<i", 9)  # after
)


def check(label, node):
    value = read_typetree(node, EndianBinaryReader(data, "<"), as_dict=True, byte_size=len(data))
    writer = EndianBinaryWriter(b"", "<")
    write_typetree(value, node, writer)
    saved = read_typetree(node, EndianBinaryReader(writer.bytes, "<"), as_dict=True, byte_size=len(writer.bytes))
    native = read_typetree(class_db, EndianBinaryReader(writer.bytes, "<"), as_dict=True, check_read=False)
    print(f"  {label}: {header(node)}")
    print(f"    read m_Script.m_PathID {value['m_Script']['m_PathID']:#x}, count/flag/after {value['count']}/{value['flag']}/{value['after']}")
    print(f"    saved bytes equal to the input: {writer.bytes == data}, m_Script.m_PathID in the saved bytes {native['m_Script']['m_PathID']:#x}")


class_db = get_typetree_node(ClassIDType.MonoBehaviour, UnityVersion.from_str(VERSION))
print("class database:", header(class_db))
for backend in ("AssetsTools", "AssetStudio", "AssetRipper"):
    generator = TypeTreeGenerator(VERSION, backend)
    generator.load_local_dll_folder(sys.argv[1])
    node = generator.get_nodes_up("Game", "Sample")
    print(f"{backend}:")
    check("generated", node)
    # the same node with the fields up to m_Name taken from the class database
    k = next(i for i, child in enumerate(node.m_Children) if child.m_Name == "m_Name") + 1
    swapped = TypeTreeNode(
        node.m_Level, node.m_Type, node.m_Name, node.m_ByteSize, node.m_Version,
        list(class_db.m_Children) + list(node.m_Children[k:]), m_MetaFlag=node.m_MetaFlag,
    )
    check("class database header", swapped)