feat: read preload table to get path id (#372)

* feat: read preload table to get path id

* refactor: only build container index once

* fix: encapsulate lock

* fix: condition

Author: StarHeartHunt

UnityPy/environment.py

@@ -32,13 +32,15 @@ class Environment:
     local_files: List[str]
     local_files_simple: List[str]
     typetree_generator: Optional["TypeTreeGenerator"] = None
+    _container_index_built: bool = False
 
     def __init__(self, *args: FileSourceType, fs: Optional[AbstractFileSystem] = None, path: Optional[str] = None):
         self.files = {}
         self.cabs = {}
         self.fs = fs or LocalFileSystem()
         self.local_files = []
         self.local_files_simple = []
+        self._container_index_built = False
 
         if path is None:
             # if no path is given, use the current working directory
@@ -203,9 +205,19 @@ def search(item):
 
         return search(self)
 
+    def _build_container_index(self) -> None:
+        if self._container_index_built:
+            return
+
+        self._container_index_built = True
+        for f in self.cabs.values():
+            if isinstance(f, SerializedFile):
+                f.container.parse_preload_table()
+
     @property
     def container(self) -> ContainerHelper:
         """Returns a dictionary of all objects in the Environment."""
+        self._build_container_index()
         container = []
         for f in self.cabs.values():
             if isinstance(f, SerializedFile) and not f.is_dependency:
@@ -249,6 +261,7 @@ def register_cab(self, name: str, item: Union[SerializedFile, EndianBinaryReader
             The file to register.
         """
         self.cabs[simplify_name(name)] = item
+        self._container_index_built = False
 
     def get_cab(self, name: str) -> Union[SerializedFile, EndianBinaryReader, None]:
         """

UnityPy/files/ObjectReader.py

@@ -198,6 +198,9 @@ def peek_name(self) -> Union[str, None]:
 
     @property
     def container(self):
+        env = self.assets_file.environment
+        if env is not None:
+            env._build_container_index()
         return self.assets_file._container.path_dict.get(self.path_id)
 
     @property

UnityPy/helpers/ContainerHelper.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Dict, Generator, Iterator, List, Tuple, Union
+from typing import TYPE_CHECKING, Dict, Generator, Iterator, List, Optional, Tuple, Union
 
 from attrs import define
 
@@ -16,13 +16,37 @@ class ContainerHelper:
     container: List[Tuple[str, AssetInfo]]
     container_dict: Dict[str, PPtr[Object]]
     path_dict: Dict[int, str]
+    _preload_table: Optional[List[PPtr[Object]]] = None
 
     def __init__(self, container: Union[List[Tuple[str, AssetInfo]], AssetBundle]) -> None:
+        preload_table: Optional[List[PPtr[Object]]] = None
         if not isinstance(container, (list)):
+            preload_table = container.m_PreloadTable
             container = container.m_Container
         self.container = container
         self.container_dict = {key: value.asset for key, value in container}
         self.path_dict = {value.asset.path_id: key for key, value in container}
+        self._preload_table = preload_table
+
+    def parse_preload_table(self) -> None:
+        if self._preload_table is None:
+            return
+
+        for path, info in self.container:
+            start = info.preloadIndex
+            size = info.preloadSize
+            if start < 0 or size <= 0 or start + size > len(self._preload_table):
+                continue
+            for pptr in self._preload_table[start : start + size]:
+                if not pptr:
+                    continue
+                try:
+                    target = pptr.deref()
+                except (FileNotFoundError, KeyError):
+                    continue
+                target.assets_file._container.path_dict.setdefault(pptr.path_id, path)
+
+        self._preload_table = None
 
     def items(self) -> Generator[Tuple[str, PPtr[Object]], None, None]:
         return ((key, value.asset) for key, value in self.container)