{{ message }}
Repository navigation
Expand file tree
/
Copy pathTypeTreeNode.py
More file actions
345 lines (287 loc) · 11.5 KB
/
Copy pathTypeTreeNode.py
File metadata and controls
345 lines (287 loc) · 11.5 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
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 ..helpers.Tpk import get_common_strings
from ..streams.EndianBinaryReader import EndianBinaryReader
from ..streams.EndianBinaryWriter import EndianBinaryWriter
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
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
CLEAN_NAME_REMOVE_RE = re.compile(r"[\?\*]")
CLEAN_NAME_REPLACE_RE = re.compile(r"[ \.:\-\[\]]")
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:]
name = CLEAN_NAME_REMOVE_RE.sub("", name)
name = CLEAN_NAME_REPLACE_RE.sub("_", name)
if name in ["pass", "from"]:
name += "_"
if name[0].isdigit():
name = f"x{name}"
return name
__all__ = (
"TypeTreeNode",
"get_common_strings",
"clean_name",
)
You can’t perform that action at this time.