from __future__ import annotations import re from struct import Struct from threading import Lock from typing import ( TYPE_CHECKING, Dict, Iterator, List, Optional, Tuple, Union, cast, ) from attrs import define, field from ..streams.EndianBinaryReader import EndianBinaryReader from ..streams.EndianBinaryWriter import EndianBinaryWriter if TYPE_CHECKING: from .Tpk import UnityVersion try: from ..UnityPyBoost import TypeTreeNode as TypeTreeNodeC # type: ignore except ImportError: @define(slots=True) class TypeTreeNodeC: m_Level: int m_Type: str m_Name: str m_ByteSize: int m_Version: int m_Children: List[TypeTreeNode] = field(factory=list) m_TypeFlags: Optional[int] = None m_VariableCount: Optional[int] = None m_Index: Optional[int] = None m_MetaFlag: Optional[int] = None m_RefTypeHash: Optional[int] = None _clean_name: str = field(init=False) def __attrs_post_init__(self): self._clean_name = clean_name(self.m_Name) def __repr__(self): return f"TypeTreeNode(m_Level={self.m_Level}, m_Type='{self.m_Type}', \ m_Name='{self.m_Name}', m_MetaFlag={self.m_MetaFlag})" TYPETREENODE_KEYS = [ "m_Level", "m_Type", "m_Name", "m_ByteSize", "m_Version", "m_Children", "m_TypeFlags", "m_VariableCount", "m_Index", "m_MetaFlag", "m_RefTypeHash", ] SYSTEM_GLOBAL_LOCK = Lock() NAME_PEEK_NODE_CACHE: dict[Tuple[str, str, int], Union[Tuple[TypeTreeNode, str], None]] = {} class TypeTreeNode(TypeTreeNodeC): def traverse(self) -> Iterator[TypeTreeNode]: stack: list[TypeTreeNode] = [self] while stack: node = stack.pop() yield node stack.extend(reversed(node.m_Children)) @classmethod def parse(cls, reader: EndianBinaryReader, version: int) -> TypeTreeNode: # stack approach is way faster than recursion # using a fake root node to avoid special case for root node dummy_node = cls(-1, "", "", 0, 0, []) dummy_root = cls(-1, "", "", 0, 0, [dummy_node]) stack: List[Tuple[TypeTreeNode, int]] = [(dummy_root, 1)] while stack: parent, count = stack[-1] if count == 1: stack.pop() else: stack[-1] = (parent, count - 1) node = cls( m_Level=parent.m_Level + 1, m_Type=reader.read_string_to_null(), m_Name=reader.read_string_to_null(), m_ByteSize=reader.read_int(), m_VariableCount=reader.read_int() if version == 2 else None, m_Index=reader.read_int() if version != 3 else None, m_TypeFlags=reader.read_int(), m_Version=reader.read_int(), m_MetaFlag=reader.read_int() if version != 3 else None, ) parent.m_Children[-count] = node children_count = reader.read_int() if children_count > 0: node.m_Children = [dummy_node] * children_count stack.append((node, children_count)) return dummy_root.m_Children[0] @classmethod def parse_blob(cls, reader: EndianBinaryReader, version: int) -> TypeTreeNode: node_count = reader.read_int() stringbuffer_size = reader.read_int() node_struct, keys = _get_blob_node_struct(reader.endian, version) struct_data = reader.read(node_struct.size * node_count) stringbuffer_reader = EndianBinaryReader(reader.read(stringbuffer_size), reader.endian) CommonString = get_common_strings() def read_string(reader: EndianBinaryReader, value: int) -> str: is_offset = (value & 0x80000000) == 0 if is_offset: reader.Position = value return reader.read_string_to_null() offset = value & 0x7FFFFFFF return CommonString.get(offset, str(offset)) fake_root: TypeTreeNode = cls(-1, "", "", 0, 0, []) stack: List[TypeTreeNode] = [fake_root] parent = fake_root prev = fake_root for raw_node in node_struct.iter_unpack(struct_data): node = cls( **dict(zip(keys[:3], raw_node[:3])), **dict(zip(keys[5:], raw_node[5:])), m_Type=read_string(stringbuffer_reader, raw_node[3]), m_Name=read_string(stringbuffer_reader, raw_node[4]), ) if node.m_Level > prev.m_Level: stack.append(parent) parent = prev elif node.m_Level < prev.m_Level: while node.m_Level <= parent.m_Level: parent = stack.pop() parent.m_Children.append(node) prev = node return fake_root.m_Children[0] @classmethod def from_list(cls, nodes: Union[List[Dict[str, Union[str, int]]], List[TypeTreeNode]]) -> TypeTreeNode: fake_root: TypeTreeNode = cls(-1, "", "", 0, 0, []) stack: List[TypeTreeNode] = [fake_root] parent = fake_root prev = fake_root # check if the nodes contain all required fields if isinstance(nodes[0], dict): if "m_Level" not in nodes[0] or "m_Type" not in nodes[0] or "m_Name" not in nodes[0]: raise ValueError("Nodes must contain at least m_Level, m_Type and m_Name") patch_dict = {} if "m_ByteSize" not in nodes[0]: patch_dict["m_ByteSize"] = 0 if "m_Version" not in nodes[0]: patch_dict["m_Version"] = 0 nodes = [cls(**node, **patch_dict) for node in nodes] # type: ignore if TYPE_CHECKING: nodes = cast(List[TypeTreeNode], nodes) for node in nodes: if node.m_Level > prev.m_Level: stack.append(parent) parent = prev elif node.m_Level < prev.m_Level: while node.m_Level <= parent.m_Level: parent = stack.pop() parent.m_Children.append(node) prev = node return fake_root.m_Children[0] def get_name_peek_node(self) -> Union[Tuple[TypeTreeNode, str], None]: global SYSTEM_GLOBAL_LOCK with SYSTEM_GLOBAL_LOCK: key = (self.m_Name, self.m_Type, self.m_Version) if key in NAME_PEEK_NODE_CACHE: return NAME_PEEK_NODE_CACHE[key] result: Union[Tuple[TypeTreeNode, str], None] = None for i, child in enumerate(self.m_Children): if child.m_Name in ("m_Name", "name"): peek_node = TypeTreeNode( self.m_Level, self.m_Type, self.m_Name, self.m_ByteSize, self.m_Version, self.m_Children[: i + 1], ) result = peek_node, child.m_Name break NAME_PEEK_NODE_CACHE[key] = result return result def dump(self, writer: EndianBinaryWriter, version: int): stack: list[TypeTreeNode] = [self] while stack: node = stack.pop() writer.write_string_to_null(self.m_Type) writer.write_string_to_null(self.m_Name) writer.write_int(self.m_ByteSize) if version == 2: assert self.m_VariableCount is not None writer.write_int(self.m_VariableCount) if version != 3: assert self.m_Index is not None writer.write_int(self.m_Index) writer.write_int(self.m_TypeFlags or 0) writer.write_int(self.m_Version) if version != 3: assert self.m_MetaFlag is not None writer.write_int(self.m_MetaFlag) writer.write_int(len(self.m_Children)) stack.extend(reversed(node.m_Children)) def dump_blob(self, writer: EndianBinaryWriter, version: int): node_writer = EndianBinaryWriter(endian=writer.endian) string_writer = EndianBinaryWriter() # string buffer setup CommonStringOffsetMap = {string: offset for offset, string in get_common_strings().items()} string_offsets: dict[str, int] = {} def write_string(string: str) -> int: offset = string_offsets.get(string) if offset is None: common_offset = CommonStringOffsetMap.get(string) if common_offset: offset = common_offset | 0x80000000 else: offset = string_writer.Position string_writer.write_string_to_null(string) string_offsets[string] = offset return offset # node buffer setup node_struct, keys = _get_blob_node_struct(writer.endian, version) def write_node(node: TypeTreeNode): node_writer.write( node_struct.pack( *[getattr(node, key) for key in keys[:3]], write_string(node.m_Type), write_string(node.m_Name), *[getattr(node, key) for key in keys[5:]], ) ) # write nodes node_count = len([write_node(node) for node in self.traverse()]) # write blob writer.write_int(node_count) writer.write_int(string_writer.Position) writer.write(node_writer.bytes) writer.write(string_writer.bytes) def dump_structure(self, indent: str = " ") -> str: # dump structure similar to https://github.com/AssetRipper/TypeTreeDumps/blob/main/StructsDump sb = [ f"{indent}{self.m_Type} {self.m_Name} // ByteSize{{{self.m_ByteSize:X}}}, Index{{{self.m_Index}}}, \ Version{{{self.m_Version}}}, TypeFlags{{{self.m_TypeFlags}}}, MetaFlag{{{self.m_MetaFlag}}}" ] for child in self.m_Children: sb.append(child.dump_structure(indent + " ")) return "\n".join(sb) def to_dict(self) -> dict: return { key: value for key, value in ((key, getattr(self, key)) for key in TYPETREENODE_KEYS) if value is not None } def to_dict_list(self) -> List[dict]: return [ self.to_dict(), *(item for child in self.m_Children for item in child.to_dict_list()), ] def __eq__(self, other: TypeTreeNode) -> bool: # type: ignore return self.to_dict() == other.to_dict() and self.m_Children == other.m_Children COMMONSTRING_CACHE: Dict[Optional[UnityVersion], Dict[int, str]] = {} def get_common_strings(version: Optional[UnityVersion] = None) -> Dict[int, str]: if version in COMMONSTRING_CACHE: return COMMONSTRING_CACHE[version] from .Tpk import TPKTYPETREE tree = TPKTYPETREE common_string = tree.CommonString strings = common_string.GetStrings(tree.StringBuffer) if version: count = common_string.GetCount(version) strings = strings[:count] ret: Dict[int, str] = {} offset = 0 for string in strings: ret[offset] = string offset += len(string) + 1 COMMONSTRING_CACHE[version] = ret return ret def _get_blob_node_struct(endian: str, version: int) -> tuple[Struct, list[str]]: struct_type = f"{endian}hBBIIiii" keys = [ "m_Version", "m_Level", "m_TypeFlags", "m_TypeStrOffset", "m_NameStrOffset", "m_ByteSize", "m_Index", "m_MetaFlag", ] if version >= 19: struct_type += "Q" keys.append("m_RefTypeHash") return Struct(struct_type), keys def clean_name(name: str) -> str: # keep in sync with TypeTreeHelper.cpp if len(name) == 0: return name if name.startswith("(int&)"): name = name[6:] if name.endswith("?"): name = name[:-1] name = re.sub(r"[ \.:\-\[\]]", "_", name) if name in ["pass", "from"]: name += "_" if name[0].isdigit(): name = f"x{name}" return name __all__ = ( "TypeTreeNode", "get_common_strings", "clean_name", )