TypeTreeGenerator impl

Author: K0lb3

README.md

@@ -8,13 +8,18 @@
 
 A Unity asset extractor for Python based on [AssetStudio](https://github.com/Perfare/AssetStudio).
 
-Next to extraction, it also supports editing Unity assets.
-So far following obj types can be edited:
-
--   Texture2D
--   Sprite(indirectly via linked Texture2D)
--   TextAsset
--   MonoBehaviour (and all other types that you have the typetree of)
+Next to extraction, UnityPy also supports editing Unity assets.
+Via the typetree structure all objects types can be edited.
+```py
+# modification via dict:
+    raw_dict = obj.read_typetree()
+    # modify raw dict
+    obj.save_typetree(raw_dict)
+# modification via parsed class
+    instance = obj.read()
+    # modify instance
+    obj.save(instance)
+```
 
 If you need advice or if you want to talk about (game) data-mining,
 feel free to join the [UnityPy Discord](https://discord.gg/C6txv7M).
@@ -270,9 +275,8 @@ for obj in env.objects:
 ### [MonoBehaviour](UnityPy/classes/MonoBehaviour.py)
 
 MonoBehaviour assets are usually used to save the class instances with their values.
-If a type tree exists, it can be used to read the whole data,
-but if it doesn't exist, then it is usually necessary to investigate the class that loads the specific MonoBehaviour to extract the data.
-([example](examples/CustomMonoBehaviour/get_scriptable_texture.py))
+The structure/typetree for these classes might not be contained in the asset files.
+In such cases see the 2nd example (TypeTreeGenerator) below.
 
 -   `.m_Name`
 -   `.m_Script`
@@ -303,6 +307,39 @@ for obj in env.objects:
                 data.save(raw_data = f.read())
 ```
 
+**TypeTreeGenerator**
+
+UnityPy can generate the typetrees of MonoBehaviours from the game assemblies using an optional package, ``TypeTreeGeneratorAPI``, which has to be installed via pip.
+UnityPy will automatically try to generate the typetree of MonoBehaviours if the typetree is missing in the assets and ``env.typetree_generator`` is set.
+
+```py
+import UnityPy
+from UnityPy.helpers.TypeTreeGenerator import TypeTreeGenerator
+
+# create generator
+GAME_ROOT_DIR: str
+# e.g. r"D:\Program Files (x86)\Steam\steamapps\common\Aethermancer Demo"
+GAME_UNITY_VERSION: str
+# you can get the version via an object
+# e.g. objects[0].assets_file.unity_version
+
+generator = TypeTreeGenerator(GAME_UNITY_VERSION)
+generator.load_local_game(GAME_ROOT_DIR)
+# generator.load_local_game(root_dir: str) - for a Windows game
+# generator.load_dll_folder(dll_dir: str) - for mono / non-il2cpp or generated dummies
+# generator.load_dll(dll: bytes)
+# generator.load_il2cpp(il2cpp: bytes, metadata: bytes)
+
+env = UnityPy.load(fp)
+# assign generator to env
+env.typetree_generator = generator
+for obj in objects:
+    if obj.type.name == "MonoBehaviour":
+        # automatically tries to use the generator in the background if necessary
+        x = obj.read()
+```
+
+
 ### [AudioClip](UnityPy/classes/AudioClip.py)
 
 -   `.samples` - `{sample-name : sample-data}`

UnityPy/environment.py

@@ -2,7 +2,7 @@
 import ntpath
 import os
 import re
-from typing import Callable, Dict, List, Optional, Union
+from typing import Callable, Dict, List, Optional, Union, TYPE_CHECKING
 from zipfile import ZipFile
 
 from fsspec import AbstractFileSystem
@@ -18,6 +18,9 @@
 )
 from .streams import EndianBinaryReader
 
+if TYPE_CHECKING:
+    from UnityPy.helpers.TypeTreeGenerator import TypeTreeGenerator
+
 reSplit = re.compile(r"(.*?([^\/\\]+?))\.split\d+")
 
 
@@ -27,6 +30,7 @@ class Environment:
     path: str
     local_files: List[str]
     local_files_simple: List[str]
+    typetree_generator: Optional["TypeTreeGenerator"] = None
 
     def __init__(self, *args: FileSourceType, fs: Optional[AbstractFileSystem] = None):
         self.files = {}

UnityPy/files/ObjectReader.py

@@ -1,8 +1,17 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Generic, List, Optional, Tuple, Type, TypeVar, Union
+from typing import (
+    TYPE_CHECKING,
+    Generic,
+    List,
+    Optional,
+    Tuple,
+    Type,
+    TypeVar,
+    Union,
+    Any,
+)
 
-from ..classes import NamedObject
 from ..classes.ClassIDTypeToClassMap import ClassIDTypeToClassMap
 from ..enums import ClassIDType, BuildTarget
 from ..exceptions import TypeTreeError
@@ -215,20 +224,20 @@ def __repr__(self):
 
     def dump_typetree_structure(
         self,
-        nodes: Optional[Union[TypeTreeNode, List[dict]]] = None,
+        nodes: Optional[Union[TypeTreeNode, List[dict[str, Union[str, int]]]]] = None,
         indent: str = "  ",
     ) -> str:
         node = self._get_typetree_node(nodes)
         return node.dump_structure(indent=indent)
 
     def read_typetree(
         self,
-        nodes: Optional[Union[TypeTreeNode, List[dict]]] = None,
+        nodes: Optional[Union[TypeTreeNode, List[dict[str, Union[str, int]]]]] = None,
         wrap: bool = False,
         check_read: bool = True,
     ) -> Union[dict, T]:
-        self.reset()
         node = self._get_typetree_node(nodes)
+        self.reset()
         ret = TypeTreeHelper.read_typetree(
             node,
             self.reader,
@@ -244,7 +253,7 @@ def read_typetree(
     def save_typetree(
         self,
         tree: dict,
-        nodes: Optional[Union[TypeTreeNode, List[dict]]] = None,
+        nodes: Optional[Union[TypeTreeNode, List[dict[str, Union[str, int]]]]] = None,
         writer: EndianBinaryWriter = None,
     ):
         node = self._get_typetree_node(nodes)
@@ -263,7 +272,8 @@ def get_raw_data(self) -> bytes:
         return ret
 
     def _get_typetree_node(
-        self, node: Optional[Union[TypeTreeNode, List[dict]]] = None
+        self,
+        node: Optional[Union[TypeTreeNode, List[dict[str, Union[str, int]]]]] = None,
     ) -> TypeTreeNode:
         if isinstance(node, TypeTreeNode):
             return node
@@ -276,13 +286,33 @@ def _get_typetree_node(
             node = self.serialized_type.node
         if not node:
             node = get_typetree_node(self.class_id, self.version)
+            if node.m_Type == "MonoBehaviour":
+                try:
+                    node = self._try_monobehaviour_node(node)
+                except Exception:
+                    pass
         if not node:
             raise TypeTreeError("There are no TypeTree nodes for this object.")
         return node
 
     # UnityPy 2 syntax early implementation
-    def parse_as_object(self) -> T:
-        return self.read()
-
-    def parse_as_dict(self) -> dict:
-        return self.read_typetree()
+    def parse_as_object(self, check_read: bool = True) -> T:
+        return self.read(check_read)
+
+    def parse_as_dict(self, check_read: bool = True) -> Union[dict[str, Any], T]:
+        return self.read_typetree(check_read=check_read)
+
+    def _try_monobehaviour_node(self, base_node: TypeTreeNode) -> TypeTreeNode:
+        env = self.assets_file.environment
+        generator = env.typetree_generator
+        if generator is None:
+            raise ValueError("No typetree generator set!")
+        monobehaviour = self.read_typetree(base_node, check_read=False, wrap=True)
+        script = monobehaviour.m_Script.deref_parse_as_object()
+        node = generator.get_nodes_up(
+            script.m_AssemblyName, f"{script.m_Namespace}.{script.m_ClassName}"
+        )
+        if node:
+            return node
+        else:
+            raise ValueError("Failed to get custom MonoBehaviour node!")

UnityPy/helpers/TypeTreeGenerator.py

@@ -0,0 +1,92 @@
+import os
+from typing import Dict, List, Tuple
+from .TypeTreeNode import TypeTreeNode
+
+try:
+    from TypeTreeGeneratorAPI import TypeTreeGenerator as TypeTreeGeneratorBase
+except ImportError:
+
+    class TypeTreeGeneratorBase:
+        def __init__(self, unity_version: str):
+            raise ImportError("TypeTreeGeneratorAPI isn't installed!")
+
+        def load_dll(self, dll: bytes): ...
+        def load_il2cpp(self, il2cpp: bytes, metadata: bytes): ...
+        def get_nodes_as_json(self, assembly: str, fullname: str) -> str: ...
+        def get_nodes(self, assembly: str, fullname: str) -> List[TypeTreeNode]: ...
+
+
+class TypeTreeGenerator(TypeTreeGeneratorBase):
+    cache: Dict[Tuple[str, str], TypeTreeNode]
+
+    def __init__(self, unity_version: str):
+        super().__init__(unity_version)
+        self.cache = {}
+
+    def load_local_game(self, root_dir: str):
+        root_files = os.listdir(root_dir)
+        data_dir = os.path.join(
+            root_dir, next(f for f in root_files if f.endswith("_Data"))
+        )
+        if "GameAssembly.dll" in root_files:
+            ga_fp = os.path.join(root_dir, "GameAssembly.dll")
+            gm_fp = os.path.join(
+                data_dir, "il2cpp_data", "Metadata", "global-metadata.dat"
+            )
+            ga_raw = open(ga_fp, "rb").read()
+            gm_raw = open(gm_fp, "rb").read()
+            self.load_il2cpp(ga_raw, gm_raw)
+        else:
+            self.load_local_dll_folder(os.path.join(data_dir, "Managed"))
+
+    def load_local_dll_folder(self, dll_dir: str):
+        for f in os.listdir(dll_dir):
+            fp = os.path.join(dll_dir, f)
+            with open(fp, "rb") as f:
+                data = f.read()
+                self.load_dll(data)
+
+    def get_nodes_up(self, assembly: str, fullname: str) -> TypeTreeNode:
+        root = self.cache.get((assembly, fullname))
+        if root is not None:
+            return root
+
+        base_nodes = self.get_nodes(f"{assembly}.dll", fullname)
+
+        base_root = base_nodes[0]
+        root = TypeTreeNode(
+            base_root.m_Level,
+            base_root.m_Type,
+            base_root.m_Name,
+            0,
+            0,
+            m_MetaFlag=base_root.m_MetaFlag,
+        )
+        stack: List[TypeTreeNode] = []
+        parent = root
+        prev = root
+
+        for base_node in base_nodes[1:]:
+            node = TypeTreeNode(
+                base_node.m_Level,
+                base_node.m_Type,
+                base_node.m_Name,
+                0,
+                0,
+                m_MetaFlag=base_node.m_MetaFlag,
+            )
+            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
+
+        self.cache[(assembly, fullname)] = root
+        return root
+
+
+__all__ = ("TypeTreeGenerator",)

pyproject.toml

@@ -54,16 +54,18 @@ dependencies = [
 ]
 dynamic = ["version"]
 
+[project.optional-dependencies]
+typetree_api = ["TypeTreeAPI"]
+full = ["UnityPy[typetree_api]"]
+tests = ["pytest", "pillow", "psutil", "UnityPy[full]"]
+
 [project.urls]
 "Homepage" = "https://github.com/K0lb3/UnityPy"
 "Bug Tracker" = "https://github.com/K0lb3/UnityPy/issues"
 
 [tool.setuptools.dynamic]
 version = { attr = "UnityPy.__version__" }
 
-[project.optional-dependencies]
-tests = ["pytest", "pillow", "psutil"]
-
 [tool.pytest.ini_options]
 testpaths = ["tests"]