Object Parsing/Typetree Reading - add check_read arg

Author: K0lb3

README.md

@@ -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())
 ```

UnityPy/files/ObjectReader.py

@@ -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
 
@@ -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)
@@ -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)

UnityPy/helpers/TypeTreeHelper.py

@@ -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.
@@ -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

UnityPyBoost/TypeTreeHelper.cpp

@@ -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");
@@ -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);
@@ -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;
@@ -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");
@@ -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

examples/MonoBehaviourFromAssembly/monobehaviour_from_assembly.py

@@ -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

tests/test_typetree.py

@@ -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:
@@ -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}"
@@ -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)
@@ -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