Skip to content

Commit

Permalink
Object Parsing/Typetree Reading - add check_read arg
Browse files Browse the repository at this point in the history
  • Loading branch information
K0lb3 committed Oct 23, 2024
1 parent c1edfae commit 2f79790
Show file tree
Hide file tree
Showing 6 changed files with 39 additions and 37 deletions.
2 changes: 1 addition & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -298,7 +298,7 @@ for obj in env.objects:
# apply modifications to the data within the tree
obj.save_typetree(tree)
else:
data = obj.read()
data = obj.read(check_read=False)
with open(os.path.join(replace_dir, data.m_Name)) as f:
data.save(raw_data = f.read())
```
Expand Down
8 changes: 5 additions & 3 deletions UnityPy/files/ObjectReader.py
Original file line number Diff line number Diff line change
Expand Up @@ -170,8 +170,8 @@ def Position(self, pos):
def reset(self):
self.reader.Position = self.byte_start

def read(self) -> T:
obj = self.read_typetree(wrap=True)
def read(self, check_read: bool = True) -> T:
obj = self.read_typetree(wrap=True, check_read=check_read)
self._read_until = self.reader.Position
return obj

Expand Down Expand Up @@ -203,6 +203,7 @@ def read_typetree(
self,
nodes: Optional[Union[TypeTreeNode, List[dict]]] = None,
wrap: bool = False,
check_read: bool = True,
) -> Union[dict, T]:
self.reset()
node = self._get_typetree_node(nodes)
Expand All @@ -211,7 +212,8 @@ def read_typetree(
self.reader,
as_dict=not wrap,
assetsfile=self.assets_file,
expected_read=self.byte_size,
byte_size=self.byte_size,
check_read=check_read,
)
if wrap:
ret.set_object_reader(self)
Expand Down
24 changes: 13 additions & 11 deletions UnityPy/helpers/TypeTreeHelper.py
Original file line number Diff line number Diff line change
Expand Up @@ -115,7 +115,8 @@ def read_typetree(
root_node: TypeTreeNode,
reader: EndianBinaryReader,
as_dict: bool = True,
expected_read: Optional[int] = None,
byte_size: Optional[int] = None,
check_read: bool = True,
assetsfile: Optional[SerializedFile] = None,
) -> Union[dict[str, Any], Object]:
"""Reads the typetree of the object contained in the reader via the node list.
Expand All @@ -132,20 +133,21 @@ def read_typetree(
dict | objects.Object
The parsed typtree
"""
if expected_read and read_typetree_boost:
data = reader.read_bytes(expected_read)
return read_typetree_boost(
bytes_read: int
if byte_size and read_typetree_boost:
data = reader.read_bytes(byte_size)
obj, bytes_read = read_typetree_boost(
data, root_node, reader.endian, as_dict, assetsfile, classes
)
else:
pos = reader.Position
config = TypeTreeConfig(as_dict, assetsfile, False)
obj = read_value(root_node, reader, config)
bytes_read = reader.Position - pos

pos = reader.Position
config = TypeTreeConfig(as_dict, assetsfile, False)
obj = read_value(root_node, reader, config)

read = reader.Position - pos
if expected_read is not None and read != expected_read:
if check_read and bytes_read != byte_size:
raise ValueError(
f"Expected to read {expected_read} bytes, but read {read} bytes"
f"Expected to read {byte_size} bytes, but only read {bytes_read} bytes"
)

return obj
Expand Down
19 changes: 6 additions & 13 deletions UnityPyBoost/TypeTreeHelper.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -582,7 +582,7 @@ TypeTreeNodeObject *get_ref_type_node(PyObject *ref_object, PyObject *assetsfile
{
if (assetsfile == Py_None)
{
PyErr_SetString(PyExc_ValueError, "No serialized file given!");
PyErr_SetString(PyExc_ValueError, "Reference Type found but no SerializedFile passed as assetsfile to read_typetree!");
return NULL;
}
PyObject *ref_types = PyObject_GetAttrString(assetsfile, "ref_types");
Expand Down Expand Up @@ -656,7 +656,7 @@ TypeTreeNodeObject *get_ref_type_node(PyObject *ref_object, PyObject *assetsfile
break;
}

bool compare_cls = PyObject_RichCompareBool(cls, m_ClassName, Py_EQ) && PyObject_RichCompareBool(ns, m_NameSpace, Py_EQ) && PyObject_RichCompareBool(asm_, m_AssemblyName, Py_EQ);
bool compare_cls = (PyUnicode_Compare(cls, m_ClassName) == 0) && (PyUnicode_Compare(ns, m_NameSpace) == 0) && (PyUnicode_Compare(asm_, m_AssemblyName) == 0);
Py_DECREF(m_ClassName);
Py_DECREF(m_NameSpace);
Py_DECREF(m_AssemblyName);
Expand Down Expand Up @@ -939,6 +939,7 @@ PyObject *read_typetree(PyObject *self, PyObject *args, PyObject *kwargs)
PyObject *node = nullptr;
int as_dict = 1;
PyObject *value = nullptr;
Py_ssize_t bytes_read = 0;
ReaderT reader;

volatile uint16_t bint = 0x0100;
Expand All @@ -965,11 +966,6 @@ PyObject *read_typetree(PyObject *self, PyObject *args, PyObject *kwargs)
config.as_dict = as_dict == 1;
if (!config.as_dict)
{
if (config.assetfile == Py_None)
{
PyErr_SetString(PyExc_ValueError, "assetsfile must be set if not as dict");
goto READ_TYPETREE_CLEANUP;
}
if (config.classes == Py_None)
{
PyErr_SetString(PyExc_ValueError, "classes must be set if not as dict");
Expand Down Expand Up @@ -1019,17 +1015,14 @@ PyObject *read_typetree(PyObject *self, PyObject *args, PyObject *kwargs)
value = read_typetree_value<false>(&reader, (TypeTreeNodeObject *)node, &config);
}

if (reader.ptr != reader.end)
{
Py_DECREF(value);
value = PyErr_Format(PyExc_ValueError, "Read %ld bytes, %ld remaining", reader.ptr - reader.start, reader.end - reader.ptr);
}
bytes_read = reader.ptr - reader.start;

READ_TYPETREE_CLEANUP:
PyBuffer_Release(&view);
Py_XDECREF(config.assetfile);
Py_XDECREF(config.classes);
return value;

return Py_BuildValue("(Nn)", value, bytes_read);
}

// TypeTreeNode impl
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -46,10 +46,10 @@ def export_monobehaviours(asset_path: str, trees: dict):
continue
for obj in env.objects:
if obj.type == "MonoBehaviour":
d = obj.read()
if obj.serialized_type and obj.serialized_type.node:
tree = obj.read_typetree()
else:
d = obj.read(check_read=False)
if not d.m_Script:
continue
# RIP, no referenced script
Expand Down
21 changes: 13 additions & 8 deletions tests/test_typetree.py
Original file line number Diff line number Diff line change
Expand Up @@ -125,6 +125,15 @@ def generate_sample_data(
return sample_values


def _test_read_typetree(node: TypeTreeNode, data: bytes, as_dict: bool):
reader = EndianBinaryReader(data, "<")
py_values = read_typetree(node, reader, as_dict=as_dict, check_read=False)
reader.Position = 0
cpp_values = read_typetree(node, reader, as_dict=as_dict, byte_size=len(data))
assert py_values == cpp_values
return py_values


@check_leak
def test_simple_nodes():
for typs, py_typ, bounds in SIMPLE_NODE_SAMPLES:
Expand All @@ -135,7 +144,7 @@ def test_simple_nodes():
writer = EndianBinaryWriter(b"", "<")
write_typetree(value, node, writer)
raw = writer.bytes
re_value = read_typetree(node, EndianBinaryReader(raw, "<"))
re_value = _test_read_typetree(node, raw, as_dict=True)
assert (
abs(value - re_value) < 1e-5
), f"Failed on {typ}: {value} != {re_value}"
Expand All @@ -158,7 +167,7 @@ def generate_list_node(item_node: TypeTreeNode):
writer = EndianBinaryWriter(b"", "<")
write_typetree(values, array_node, writer)
raw = writer.bytes
re_values = read_typetree(array_node, EndianBinaryReader(raw, "<"))
re_values = _test_read_typetree(array_node, raw, as_dict=True)
assert all(
(abs(value - re_value) < 1e-5)
for value, re_value in zip(values, re_values)
Expand All @@ -176,19 +185,15 @@ def test_class_node_dict():
writer = EndianBinaryWriter(b"", "<")
write_typetree(TEST_CLASS_NODE_DICT, TEST_CLASS_NODE, writer)
raw = writer.bytes
re_value = read_typetree(
TEST_CLASS_NODE, EndianBinaryReader(raw, "<"), as_dict=True
)
re_value = _test_read_typetree(TEST_CLASS_NODE, raw, as_dict=True)
assert re_value == TEST_CLASS_NODE_DICT


def test_class_node_clz():
writer = EndianBinaryWriter(b"", "<")
write_typetree(TEST_CLASS_NODE_OBJ, TEST_CLASS_NODE, writer)
raw = writer.bytes
re_value = read_typetree(
TEST_CLASS_NODE, EndianBinaryReader(raw, "<"), as_dict=False
)
re_value = _test_read_typetree(TEST_CLASS_NODE, raw, as_dict=False)
assert re_value == TEST_CLASS_NODE_OBJ


Expand Down

0 comments on commit 2f79790

Please sign in to comment.