Skip to content

Commit 2f79790

Browse files
K0lb3K0lb3
committed
Object Parsing/Typetree Reading - add check_read arg
1 parent c1edfae commit 2f79790

6 files changed

Lines changed: 39 additions & 37 deletions

File tree

‎README.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -298,7 +298,7 @@ for obj in env.objects:
298298
# apply modifications to the data within the tree
299299
obj.save_typetree(tree)
300300
else:
301-
data = obj.read()
301+
data = obj.read(check_read=False)
302302
with open(os.path.join(replace_dir, data.m_Name)) as f:
303303
data.save(raw_data = f.read())
304304
```

‎UnityPy/files/ObjectReader.py‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -170,8 +170,8 @@ def Position(self, pos):
170170
def reset(self):
171171
self.reader.Position = self.byte_start
172172

173-
def read(self) -> T:
174-
obj = self.read_typetree(wrap=True)
173+
def read(self, check_read: bool = True) -> T:
174+
obj = self.read_typetree(wrap=True, check_read=check_read)
175175
self._read_until = self.reader.Position
176176
return obj
177177

@@ -203,6 +203,7 @@ def read_typetree(
203203
self,
204204
nodes: Optional[Union[TypeTreeNode, List[dict]]] = None,
205205
wrap: bool = False,
206+
check_read: bool = True,
206207
) -> Union[dict, T]:
207208
self.reset()
208209
node = self._get_typetree_node(nodes)
@@ -211,7 +212,8 @@ def read_typetree(
211212
self.reader,
212213
as_dict=not wrap,
213214
assetsfile=self.assets_file,
214-
expected_read=self.byte_size,
215+
byte_size=self.byte_size,
216+
check_read=check_read,
215217
)
216218
if wrap:
217219
ret.set_object_reader(self)

‎UnityPy/helpers/TypeTreeHelper.py‎

Lines changed: 13 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -115,7 +115,8 @@ def read_typetree(
115115
root_node: TypeTreeNode,
116116
reader: EndianBinaryReader,
117117
as_dict: bool = True,
118-
expected_read: Optional[int] = None,
118+
byte_size: Optional[int] = None,
119+
check_read: bool = True,
119120
assetsfile: Optional[SerializedFile] = None,
120121
) -> Union[dict[str, Any], Object]:
121122
"""Reads the typetree of the object contained in the reader via the node list.
@@ -132,20 +133,21 @@ def read_typetree(
132133
dict | objects.Object
133134
The parsed typtree
134135
"""
135-
if expected_read and read_typetree_boost:
136-
data = reader.read_bytes(expected_read)
137-
return read_typetree_boost(
136+
bytes_read: int
137+
if byte_size and read_typetree_boost:
138+
data = reader.read_bytes(byte_size)
139+
obj, bytes_read = read_typetree_boost(
138140
data, root_node, reader.endian, as_dict, assetsfile, classes
139141
)
142+
else:
143+
pos = reader.Position
144+
config = TypeTreeConfig(as_dict, assetsfile, False)
145+
obj = read_value(root_node, reader, config)
146+
bytes_read = reader.Position - pos
140147

141-
pos = reader.Position
142-
config = TypeTreeConfig(as_dict, assetsfile, False)
143-
obj = read_value(root_node, reader, config)
144-
145-
read = reader.Position - pos
146-
if expected_read is not None and read != expected_read:
148+
if check_read and bytes_read != byte_size:
147149
raise ValueError(
148-
f"Expected to read {expected_read} bytes, but read {read} bytes"
150+
f"Expected to read {byte_size} bytes, but only read {bytes_read} bytes"
149151
)
150152

151153
return obj

‎UnityPyBoost/TypeTreeHelper.cpp‎

Lines changed: 6 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -582,7 +582,7 @@ TypeTreeNodeObject *get_ref_type_node(PyObject *ref_object, PyObject *assetsfile
582582
{
583583
if (assetsfile == Py_None)
584584
{
585-
PyErr_SetString(PyExc_ValueError, "No serialized file given!");
585+
PyErr_SetString(PyExc_ValueError, "Reference Type found but no SerializedFile passed as assetsfile to read_typetree!");
586586
return NULL;
587587
}
588588
PyObject *ref_types = PyObject_GetAttrString(assetsfile, "ref_types");
@@ -656,7 +656,7 @@ TypeTreeNodeObject *get_ref_type_node(PyObject *ref_object, PyObject *assetsfile
656656
break;
657657
}
658658

659-
bool compare_cls = PyObject_RichCompareBool(cls, m_ClassName, Py_EQ) && PyObject_RichCompareBool(ns, m_NameSpace, Py_EQ) && PyObject_RichCompareBool(asm_, m_AssemblyName, Py_EQ);
659+
bool compare_cls = (PyUnicode_Compare(cls, m_ClassName) == 0) && (PyUnicode_Compare(ns, m_NameSpace) == 0) && (PyUnicode_Compare(asm_, m_AssemblyName) == 0);
660660
Py_DECREF(m_ClassName);
661661
Py_DECREF(m_NameSpace);
662662
Py_DECREF(m_AssemblyName);
@@ -939,6 +939,7 @@ PyObject *read_typetree(PyObject *self, PyObject *args, PyObject *kwargs)
939939
PyObject *node = nullptr;
940940
int as_dict = 1;
941941
PyObject *value = nullptr;
942+
Py_ssize_t bytes_read = 0;
942943
ReaderT reader;
943944

944945
volatile uint16_t bint = 0x0100;
@@ -965,11 +966,6 @@ PyObject *read_typetree(PyObject *self, PyObject *args, PyObject *kwargs)
965966
config.as_dict = as_dict == 1;
966967
if (!config.as_dict)
967968
{
968-
if (config.assetfile == Py_None)
969-
{
970-
PyErr_SetString(PyExc_ValueError, "assetsfile must be set if not as dict");
971-
goto READ_TYPETREE_CLEANUP;
972-
}
973969
if (config.classes == Py_None)
974970
{
975971
PyErr_SetString(PyExc_ValueError, "classes must be set if not as dict");
@@ -1019,17 +1015,14 @@ PyObject *read_typetree(PyObject *self, PyObject *args, PyObject *kwargs)
10191015
value = read_typetree_value<false>(&reader, (TypeTreeNodeObject *)node, &config);
10201016
}
10211017

1022-
if (reader.ptr != reader.end)
1023-
{
1024-
Py_DECREF(value);
1025-
value = PyErr_Format(PyExc_ValueError, "Read %ld bytes, %ld remaining", reader.ptr - reader.start, reader.end - reader.ptr);
1026-
}
1018+
bytes_read = reader.ptr - reader.start;
10271019

10281020
READ_TYPETREE_CLEANUP:
10291021
PyBuffer_Release(&view);
10301022
Py_XDECREF(config.assetfile);
10311023
Py_XDECREF(config.classes);
1032-
return value;
1024+
1025+
return Py_BuildValue("(Nn)", value, bytes_read);
10331026
}
10341027

10351028
// TypeTreeNode impl

‎examples/MonoBehaviourFromAssembly/monobehaviour_from_assembly.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,10 +46,10 @@ def export_monobehaviours(asset_path: str, trees: dict):
4646
continue
4747
for obj in env.objects:
4848
if obj.type == "MonoBehaviour":
49-
d = obj.read()
5049
if obj.serialized_type and obj.serialized_type.node:
5150
tree = obj.read_typetree()
5251
else:
52+
d = obj.read(check_read=False)
5353
if not d.m_Script:
5454
continue
5555
# RIP, no referenced script

‎tests/test_typetree.py‎

Lines changed: 13 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -125,6 +125,15 @@ def generate_sample_data(
125125
return sample_values
126126

127127

128+
def _test_read_typetree(node: TypeTreeNode, data: bytes, as_dict: bool):
129+
reader = EndianBinaryReader(data, "<")
130+
py_values = read_typetree(node, reader, as_dict=as_dict, check_read=False)
131+
reader.Position = 0
132+
cpp_values = read_typetree(node, reader, as_dict=as_dict, byte_size=len(data))
133+
assert py_values == cpp_values
134+
return py_values
135+
136+
128137
@check_leak
129138
def test_simple_nodes():
130139
for typs, py_typ, bounds in SIMPLE_NODE_SAMPLES:
@@ -135,7 +144,7 @@ def test_simple_nodes():
135144
writer = EndianBinaryWriter(b"", "<")
136145
write_typetree(value, node, writer)
137146
raw = writer.bytes
138-
re_value = read_typetree(node, EndianBinaryReader(raw, "<"))
147+
re_value = _test_read_typetree(node, raw, as_dict=True)
139148
assert (
140149
abs(value - re_value) < 1e-5
141150
), f"Failed on {typ}: {value} != {re_value}"
@@ -158,7 +167,7 @@ def generate_list_node(item_node: TypeTreeNode):
158167
writer = EndianBinaryWriter(b"", "<")
159168
write_typetree(values, array_node, writer)
160169
raw = writer.bytes
161-
re_values = read_typetree(array_node, EndianBinaryReader(raw, "<"))
170+
re_values = _test_read_typetree(array_node, raw, as_dict=True)
162171
assert all(
163172
(abs(value - re_value) < 1e-5)
164173
for value, re_value in zip(values, re_values)
@@ -176,19 +185,15 @@ def test_class_node_dict():
176185
writer = EndianBinaryWriter(b"", "<")
177186
write_typetree(TEST_CLASS_NODE_DICT, TEST_CLASS_NODE, writer)
178187
raw = writer.bytes
179-
re_value = read_typetree(
180-
TEST_CLASS_NODE, EndianBinaryReader(raw, "<"), as_dict=True
181-
)
188+
re_value = _test_read_typetree(TEST_CLASS_NODE, raw, as_dict=True)
182189
assert re_value == TEST_CLASS_NODE_DICT
183190

184191

185192
def test_class_node_clz():
186193
writer = EndianBinaryWriter(b"", "<")
187194
write_typetree(TEST_CLASS_NODE_OBJ, TEST_CLASS_NODE, writer)
188195
raw = writer.bytes
189-
re_value = read_typetree(
190-
TEST_CLASS_NODE, EndianBinaryReader(raw, "<"), as_dict=False
191-
)
196+
re_value = _test_read_typetree(TEST_CLASS_NODE, raw, as_dict=False)
192197
assert re_value == TEST_CLASS_NODE_OBJ
193198

194199

0 commit comments

Comments
 (0)