-
-
Notifications
You must be signed in to change notification settings - Fork 176
Expand file tree
/
Copy pathObjectReader.py
More file actions
352 lines (293 loc) · 10.9 KB
/
ObjectReader.py
File metadata and controls
352 lines (293 loc) · 10.9 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
from __future__ import annotations
from typing import (
TYPE_CHECKING,
Any,
Dict,
Generic,
List,
Optional,
Type,
TypeVar,
Union,
cast,
)
from attrs import define
from ..classes import MonoBehaviour
from ..classes.ClassIDTypeToClassMap import ClassIDTypeToClassMap
from ..enums import ClassIDType
from ..exceptions import TypeTreeError
from ..helpers import TypeTreeHelper
from ..helpers.Tpk import get_typetree_node
from ..helpers.TypeTreeNode import TypeTreeNode
from ..streams import EndianBinaryReader, EndianBinaryWriter
if TYPE_CHECKING:
from ..files.SerializedFile import SerializedFile, SerializedType
T = TypeVar("T")
NodeInput = Union[TypeTreeNode, List[Dict[str, Union[str, int]]]]
@define(
slots=True,
)
class ObjectReader(Generic[T]):
assets_file: SerializedFile
reader: EndianBinaryReader
path_id: int
type_id: int
serialized_type: Optional[SerializedType]
class_id: int
type: ClassIDType
byte_start: int
byte_size: int
is_destroyed: Optional[int]
is_stripped: Optional[int]
data: Optional[bytes] = None
_read_until: Optional[int] = None
@property
def version(self):
return self.assets_file.version
@property
def version2(self):
return self.assets_file.header.version
@property
def platform(self):
return self.assets_file.target_platform
# saves where the parser stopped
# in case that not all data is read
# and the obj.data is changed, the unknown data can be added again
@classmethod
def from_reader(cls, assets_file: SerializedFile, reader: EndianBinaryReader) -> ObjectReader[Any]:
header = assets_file.header
types = assets_file.types
# AssetStudio ObjectInfo init
if assets_file.big_id_enabled:
path_id = reader.read_long()
elif header.version < 14:
path_id = reader.read_int()
else:
reader.align_stream()
path_id = reader.read_long()
if header.version >= 22:
byte_start = reader.read_long()
else:
byte_start = reader.read_u_int()
byte_start += header.data_offset
byte_size = reader.read_u_int()
type_id = reader.read_int()
serialized_type = None
if header.version < 16:
class_id = reader.read_u_short()
for typ in types:
if typ.class_id == type_id:
serialized_type = typ
break
else:
typ = types[type_id]
serialized_type = typ
class_id = typ.class_id
clz_type = ClassIDType(class_id)
is_destroyed = None
if header.version < 11:
is_destroyed = reader.read_u_short()
if 11 <= header.version < 17:
script_type_index = reader.read_short()
if serialized_type:
serialized_type.script_type_index = script_type_index
is_stripped = None
if header.version == 15 or header.version == 16:
is_stripped = reader.read_byte()
return cls(
assets_file,
reader,
path_id,
type_id,
serialized_type,
class_id,
clz_type,
byte_start,
byte_size,
is_destroyed,
is_stripped,
)
def write(self, header, writer: EndianBinaryWriter, data_writer: EndianBinaryWriter):
if self.assets_file.big_id_enabled:
writer.write_long(self.path_id)
elif header.version < 14:
writer.write_int(self.path_id)
else:
writer.align_stream()
writer.write_long(self.path_id)
if self.data is not None:
data = self.data
# in some cases the parser doesn't read all of the object data
# games might still require the missing data
# so following code appends the missing data back to edited objects
# -> following solution has let to some problems, so it will be removed for now
# if self.type != ClassIDType.MonoBehaviour:
# end_pos = self.byte_start + self.byte_size
# if self._read_until and self._read_until != end_pos:
# self.reader.Position = self._read_until
# data += self.reader.read_bytes(end_pos - self._read_until)
else:
self.reset()
data = self.reader.read(self.byte_size)
if header.version >= 22:
writer.write_long(data_writer.Position)
else:
writer.write_u_int(data_writer.Position)
writer.write_u_int(len(data))
data_writer.write(data)
writer.write_int(self.type_id)
if header.version < 16:
writer.write_u_short(self.class_id)
if header.version < 11:
assert self.is_destroyed is not None
writer.write_u_short(self.is_destroyed)
if 11 <= header.version < 17:
assert self.serialized_type is not None
writer.write_short(self.serialized_type.script_type_index)
if header.version == 15 or header.version == 16:
assert self.is_stripped is not None
writer.write_byte(self.is_stripped)
def set_raw_data(self, data: bytes):
self.data = data
if self.assets_file:
self.assets_file.mark_changed()
def get_class(self) -> Union[Type[T], None]:
return ClassIDTypeToClassMap.get(self.type) # type: ignore
def peek_name(self) -> Union[str, None]:
"""Peeks the name of the object without reading/parsing the whole object."""
node = self._get_typetree_node()
peek_node = node.get_name_peek_node()
if peek_node:
node, key = peek_node
return self.parse_as_dict(node, check_read=False)[key]
else:
return None
@property
def container(self):
return self.assets_file._container.path_dict.get(self.path_id)
@property
def Position(self):
return self.reader.Position
@Position.setter
def Position(self, pos):
self.reader.Position = pos
def reset(self):
self.reader.Position = self.byte_start
def read(self, check_read: bool = True) -> T:
return self.read_typetree(wrap=True, check_read=check_read) # type: ignore
def get(self, key, default=None):
return getattr(self, key, default)
def __getattr__(self, item: str):
if hasattr(self.reader, item):
return getattr(self.reader, item)
def __repr__(self):
return "<%s %s>" % (self.__class__.__name__, self.type.name)
###################################################
#
# Typetree Stuff
#
###################################################
def dump_typetree_structure(
self,
nodes: Optional[NodeInput] = None,
indent: str = " ",
) -> str:
node = self._get_typetree_node(nodes)
return node.dump_structure(indent=indent)
def read_typetree(
self,
nodes: Optional[NodeInput] = None,
wrap: bool = False,
check_read: bool = True,
) -> Union[dict, T]:
node = self._get_typetree_node(nodes)
self.reset()
ret = TypeTreeHelper.read_typetree(
node,
self.reader,
as_dict=not wrap,
assetsfile=self.assets_file,
byte_size=self.byte_size,
check_read=check_read,
)
if wrap:
ret.set_object_reader(self) # type: ignore
return ret # type: ignore
def save_typetree(
self,
tree: Union[dict, T],
nodes: Optional[NodeInput] = None,
writer: Optional[EndianBinaryWriter] = None,
):
node = self._get_typetree_node(nodes)
if not writer:
writer = EndianBinaryWriter(endian=self.reader.endian)
TypeTreeHelper.write_typetree(tree, node, writer, self.assets_file)
data = writer.bytes
self.set_raw_data(data)
return data
def get_raw_data(self) -> bytes:
pos = self.Position
self.reset()
ret = self.reader.read_bytes(self.byte_size)
self.Position = pos
return ret
def _get_typetree_node(
self,
node: Optional[NodeInput] = None,
) -> TypeTreeNode:
if isinstance(node, TypeTreeNode):
return node
elif isinstance(node, list):
return TypeTreeNode.from_list(node)
elif node is not None:
raise ValueError("nodes must be a list[dict] or TypeTreeNode")
if self.serialized_type:
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.generate_monobehaviour_node(node)
except ValueError:
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, node: Optional[NodeInput] = None, check_read: bool = True) -> T:
return self.read_typetree(nodes=node, wrap=True, check_read=check_read) # type: ignore
def parse_as_dict(self, node: Optional[NodeInput] = None, check_read: bool = True) -> dict[str, Any]:
return self.read_typetree(nodes=node, wrap=False, check_read=check_read) # type: ignore
def patch(
self,
obj: Union[dict, T],
nodes: Optional[NodeInput] = None,
writer: Optional[EndianBinaryWriter] = None,
):
return self.save_typetree(obj, nodes=nodes, writer=writer)
# MonoBehaviour specific methods
def parse_monobehaviour_head(self, mb_node: Optional[TypeTreeNode] = None) -> MonoBehaviour:
if mb_node is None:
mb_node = get_typetree_node(ClassIDType.MonoBehaviour, self.version)
mb = self.read_typetree(nodes=mb_node, wrap=True, check_read=False)
return cast(MonoBehaviour, mb)
def generate_monobehaviour_node(self, mb_node: Optional[TypeTreeNode] = None) -> TypeTreeNode:
env = self.assets_file.environment
generator = env.typetree_generator
if generator is None:
raise ValueError("MonoBehaviour detected, but no typetree_generator set to the environment!")
monobehaviour = self.parse_monobehaviour_head(mb_node)
script = monobehaviour.m_Script.deref_parse_as_object()
if script.m_Namespace != "":
fullname = f"{script.m_Namespace}.{script.m_ClassName}"
else:
fullname = script.m_ClassName
node = generator.get_nodes_up(script.m_AssemblyName, fullname)
if node:
return node
else:
raise ValueError(f"Failed to generate MonoBehaviour node for {fullname} of {script.m_AssemblyName}!")
__all__ = [
"ObjectReader",
]