Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 13 additions & 0 deletions UnityPy/environment.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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]:
"""
Expand Down
3 changes: 3 additions & 0 deletions UnityPy/files/ObjectReader.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
26 changes: 25 additions & 1 deletion UnityPy/helpers/ContainerHelper.py
Original file line number Diff line number Diff line change
@@ -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

Expand All @@ -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)
Expand Down
Loading