Skip to content

Commit 52ce553

Browse files
K0lb3K0lb3
committed
TypeTree - implement ManagedReferencesRegistry reading
1 parent 40487fc commit 52ce553

4 files changed

Lines changed: 262 additions & 45 deletions

File tree

‎UnityPy/files/SerializedFile.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -214,6 +214,7 @@ class SerializedFile(File.File):
214214
types: List[SerializedType]
215215
script_types: List[LocalSerializedObjectIdentifier]
216216
externals: List[FileIdentifier]
217+
ref_types: Optional[List[SerializedType]]
217218
objects: Dict[int, ObjectReader]
218219
unknown: int
219220
header: SerializedFileHeader

‎UnityPy/helpers/TypeTreeHelper.py‎

Lines changed: 87 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,15 @@
11
from __future__ import annotations
2+
23
import re
3-
from typing import Optional, Any, Union, TYPE_CHECKING
4+
from copy import copy
5+
from typing import TYPE_CHECKING, Any, Optional, Union
46

5-
from .TypeTreeNode import TypeTreeNode
6-
from ..streams.EndianBinaryReader import EndianBinaryReader
7-
from ..streams.EndianBinaryWriter import EndianBinaryWriter
7+
from attrs import define
88

99
from .. import classes
10+
from ..streams.EndianBinaryReader import EndianBinaryReader
11+
from ..streams.EndianBinaryWriter import EndianBinaryWriter
12+
from .TypeTreeNode import TypeTreeNode
1013

1114
Object = classes.Object
1215
PPtr = classes.PPtr
@@ -70,6 +73,38 @@
7073
}
7174

7275

76+
@define(slots=True)
77+
class TypeTreeConfig:
78+
as_dict: bool
79+
assetsfile: "Optional[SerializedFile]" = None
80+
has_registry: bool = False
81+
82+
83+
def get_ref_type_node(ref_object: dict, assetfile: SerializedFile) -> TypeTreeNode:
84+
typ = ref_object["type"]
85+
if isinstance(typ, dict):
86+
cls = typ["class"]
87+
ns = typ["ns"]
88+
asm = typ["asm"]
89+
else:
90+
cls = getattr(typ, "class")
91+
ns = typ.ns
92+
asm = typ.asm
93+
94+
if not assetfile or not assetfile.ref_types:
95+
raise ValueError("SerializedFile has no ref_types")
96+
97+
for ref_type in assetfile.ref_types:
98+
if (
99+
cls == ref_type.m_ClassName
100+
and ns == ref_type.m_NameSpace
101+
and asm == ref_type.m_AssemblyName
102+
):
103+
return ref_type.node
104+
else:
105+
raise ValueError(f"Referenced type not found: {cls} {ns} {asm}")
106+
107+
73108
def read_typetree(
74109
root_node: TypeTreeNode,
75110
reader: EndianBinaryReader,
@@ -98,7 +133,8 @@ def read_typetree(
98133
)
99134

100135
pos = reader.Position
101-
obj = read_value(root_node, reader, as_dict, assetsfile)
136+
config = TypeTreeConfig(as_dict, assetsfile, False)
137+
obj = read_value(root_node, reader, config)
102138

103139
read = reader.Position - pos
104140
if expected_read is not None and read != expected_read:
@@ -131,8 +167,7 @@ def write_typetree(
131167
def read_value(
132168
node: TypeTreeNode,
133169
reader: EndianBinaryReader,
134-
as_dict: bool,
135-
assetsfile: Optional[SerializedFile],
170+
config: TypeTreeConfig,
136171
) -> Any:
137172
# print(reader.Position, node.m_Name, node.m_Type, node.m_MetaFlag)
138173
align = metaflag_is_aligned(node.m_MetaFlag)
@@ -141,9 +176,17 @@ def read_value(
141176
if func:
142177
value = func(reader)
143178
elif node.m_Type == "pair":
144-
first = read_value(node.m_Children[0], reader, as_dict, assetsfile)
145-
second = read_value(node.m_Children[1], reader, as_dict, assetsfile)
179+
first = read_value(node.m_Children[0], reader, config)
180+
second = read_value(node.m_Children[1], reader, config)
146181
value = (first, second)
182+
elif node.m_Type == "ReferencedObject":
183+
value = {}
184+
for child in node.m_Children:
185+
if child.m_Type == "ReferencedObjectData":
186+
ref_type_nodes = get_ref_type_node(value, config.assetsfile)
187+
value[child.m_Name] = read_value(ref_type_nodes, reader, config)
188+
else:
189+
value[child.m_Name] = read_value(child, reader, config)
147190
# Vector
148191
elif node.m_Children and node.m_Children[0].m_Type == "Array":
149192
if metaflag_is_aligned(node.m_Children[0].m_MetaFlag):
@@ -153,22 +196,25 @@ def read_value(
153196
size = reader.read_int()
154197
subtype = node.m_Children[0].m_Children[1]
155198
if metaflag_is_aligned(subtype.m_MetaFlag):
156-
value = read_value_array(subtype, reader, as_dict, size, assetsfile)
199+
value = read_value_array(subtype, reader, config)
157200
else:
158-
value = [
159-
read_value(subtype, reader, as_dict, assetsfile) for _ in range(size)
160-
]
201+
value = [read_value(subtype, reader, config) for _ in range(size)]
161202

162203
else: # Class
163-
value = {
164-
child.m_Name: read_value(child, reader, as_dict, assetsfile)
165-
for child in node.m_Children
166-
}
204+
value = {}
205+
for child in node.m_Children:
206+
if child.m_Type == "ManagedReferencesRegistry":
207+
if config.has_registry:
208+
continue
209+
else:
210+
config = copy(config)
211+
config.has_registry = True
212+
value[child.m_Name] = read_value(child, reader, config)
167213

168-
if not as_dict:
214+
if not config.as_dict:
169215
if node.m_Type.startswith("PPtr<"):
170216
value = PPtr[Any](
171-
assetsfile=assetsfile,
217+
assetsfile=config.assetsfile,
172218
m_FileID=value["m_FileID"],
173219
m_PathID=value["m_PathID"],
174220
)
@@ -198,9 +244,8 @@ def read_value(
198244
def read_value_array(
199245
node: TypeTreeNode,
200246
reader: EndianBinaryReader,
201-
as_dict: bool,
247+
config: TypeTreeConfig,
202248
size: int,
203-
assetsfile: Optional[SerializedFile],
204249
) -> Any:
205250
align = metaflag_is_aligned(node.m_MetaFlag)
206251

@@ -217,38 +262,44 @@ def read_value_array(
217262

218263
key_func = FUNCTION_READ_MAP.get(
219264
key_node.m_Type,
220-
lambda reader: read_value(key_node, reader, as_dict, assetsfile),
265+
lambda reader: read_value(key_node, reader, config),
221266
)
222267
value_func = FUNCTION_READ_MAP.get(
223268
value_node.m_Type,
224-
lambda reader: read_value(value_node, reader, as_dict, assetsfile),
269+
lambda reader: read_value(value_node, reader, config),
225270
)
226271
value = [(key_func(reader), value_func(reader)) for _ in range(size)]
272+
elif node.m_Type == "ReferencedObject":
273+
value = [None] * size
274+
for i in range(size):
275+
item = {}
276+
for child in node.m_Children:
277+
if child.m_Type == "ReferencedObjectData":
278+
ref_type_nodes = get_ref_type_node(item, config.assetsfile)
279+
item[child.m_Name] = read_value(ref_type_nodes, reader, config)
280+
else:
281+
item[child.m_Name] = read_value(child, reader, config)
282+
value[i] = item
227283
# Vector
228284
elif node.m_Children and node.m_Children[0].m_Type == "Array":
229285
if metaflag_is_aligned(node.m_Children[0].m_MetaFlag):
230286
align = True
231287
subtype = node.m_Children[0].m_Children[1]
232288
if metaflag_is_aligned(subtype.m_MetaFlag):
233289
value = [
234-
read_value_array(
235-
subtype, reader, as_dict, reader.read_int(), assetsfile
236-
)
290+
read_value_array(subtype, reader, config, reader.read_int())
237291
for _ in range(size)
238292
]
239293
else:
240294
value = [
241-
[
242-
read_value(subtype, reader, as_dict, assetsfile)
243-
for _ in range(reader.read_int())
244-
]
295+
[read_value(subtype, reader, config) for _ in range(reader.read_int())]
245296
for _ in range(size)
246297
]
247298
else: # Class
248-
if as_dict:
299+
if config.as_dict:
249300
value = [
250301
{
251-
child.m_Name: read_value(child, reader, as_dict, assetsfile)
302+
child.m_Name: read_value(child, reader, config)
252303
for child in node.m_Children
253304
}
254305
for _ in range(size)
@@ -257,9 +308,9 @@ def read_value_array(
257308
if node.m_Type.startswith("PPtr<"):
258309
value = [
259310
PPtr[Any](
260-
assetsfile=assetsfile,
311+
assetsfile=config.assetsfile,
261312
**{
262-
child.m_Name: read_value(child, reader, as_dict, assetsfile)
313+
child.m_Name: read_value(child, reader, config)
263314
for child in node.m_Children
264315
},
265316
)
@@ -276,7 +327,7 @@ def read_value_array(
276327
value = [
277328
clz(
278329
**{
279-
name: read_value(child, reader, as_dict, assetsfile)
330+
name: read_value(child, reader, config)
280331
for name, child in zip(clean_names, node.m_Children)
281332
}
282333
)
@@ -287,9 +338,7 @@ def read_value_array(
287338
value = [None] * size
288339
for i in range(size):
289340
value_i_d = {
290-
clean_name(child.m_Name): read_value(
291-
child, reader, as_dict, assetsfile
292-
)
341+
clean_name(child.m_Name): read_value(child, reader, config)
293342
for child in node.m_Children
294343
}
295344
value_i = clz(

0 commit comments

Comments
 (0)