Skip to content

Commit bbde69e

Browse files
K0lb3K0lb3
committed
add improved AssetBatchConverter example as tools to the lib
1 parent d61dfac commit bbde69e

2 files changed

Lines changed: 324 additions & 179 deletions

File tree

‎UnityPy/tools/extractor.py‎

Lines changed: 324 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,324 @@
1+
from io import BytesIO
2+
import os
3+
import json
4+
import UnityPy
5+
from UnityPy.classes import (
6+
Object,
7+
PPtr,
8+
MonoBehaviour,
9+
TextAsset,
10+
Font,
11+
Shader,
12+
Mesh,
13+
Sprite,
14+
Texture2D,
15+
AudioClip,
16+
GameObject,
17+
)
18+
from UnityPy.enums.ClassIDType import ClassIDType
19+
from typing import Union, List, Dict
20+
from pathlib import Path
21+
22+
23+
def export_obj(
24+
obj: Union[Object, PPtr],
25+
fp: Path,
26+
append_name: bool = False,
27+
append_path_id: bool = False,
28+
export_unknown_as_typetree: bool = False,
29+
) -> List[int]:
30+
"""Exports the given object to the given filepath.
31+
32+
Args:
33+
obj (Object, PPtr): A valid Unity object or a reference to one.
34+
fp (Path): A valid filepath where the object should be exported to.
35+
append_name (bool, optional): Decides if the obj name will be appended to the filepath. Defaults to False.
36+
append_path_id (bool, optional): Decides if the obj name will be appended to the filepath. Defaults to False.
37+
export_unknown_as_typetree (bool, optional): If set, then unimplemented objects will be exported via their typetree or dumped as bin. Defaults to False.
38+
39+
Returns:
40+
list: a list of exported object path_ids
41+
"""
42+
# figure out export function
43+
type_name = obj.type.name
44+
export_func = getattr(EXPORT_TYPES, type_name)
45+
if export_unknown_as_typetree:
46+
export_func = exportMonoBehaviour
47+
else:
48+
return []
49+
50+
# set filepath
51+
obj = obj.read()
52+
53+
if append_name:
54+
fp = os.path.join(fp, obj.name if obj.name else type_name)
55+
56+
fp, extension = os.path.splitext(fp)
57+
58+
if append_path_id:
59+
fp = f"{fp}_{obj.path_id}"
60+
61+
# export
62+
return export_func(obj, fp, extension)
63+
64+
65+
def extract_assets(
66+
src: Union[Path, BytesIO, bytes, bytearray],
67+
dst: Path,
68+
use_container: bool = True,
69+
ignore_first_container_dirs: int = 0,
70+
append_path_id: bool = False,
71+
export_unknown_as_typetree: bool = False,
72+
) -> List[int]:
73+
"""Extracts all assets from the given source.
74+
75+
Args:
76+
src (Union[Path, BytesIO, bytes, bytearray]): [description]
77+
dst (Path): [description]
78+
use_container (bool, optional): [description]. Defaults to True.
79+
ignore_first_container_dirs (int, optional): [description]. Defaults to 0.
80+
append_path_id (bool, optional): [description]. Defaults to False.
81+
export_unknown_as_typetree (bool, optional): [description]. Defaults to False.
82+
83+
Returns:
84+
List[int]: [description]
85+
"""
86+
# load source
87+
env = UnityPy.load(src)
88+
exported = []
89+
90+
export_types_keys = list(EXPORT_TYPES.keys())
91+
92+
def defaulted_export_index(type: ClassIDType):
93+
try:
94+
return export_types_keys.index(type)
95+
except IndexError:
96+
return 999
97+
98+
if use_container:
99+
container = sorted(env.container, lambda x: defaulted_export_index(x[1].type))
100+
for obj_path, obj in container:
101+
# the check of the various sub directories is required to avoid // in the path
102+
obj_dest = os.path.join(
103+
dst,
104+
*(x for x in obj_path.split("/")[:ignore_first_container_dirs] if x),
105+
)
106+
os.makedirs(os.path.dirname(obj_dest), exist_ok=True)
107+
exported.extend(
108+
export_obj(
109+
obj,
110+
obj_dest,
111+
append_path_id=append_path_id,
112+
export_unknown_as_typetree=export_unknown_as_typetree,
113+
)
114+
)
115+
116+
else:
117+
objects = sorted(env.objects, lambda x: defaulted_export_index(x.type))
118+
for obj in objects:
119+
if obj.path_id not in exported:
120+
exported.extend(
121+
export_obj(
122+
obj,
123+
dst,
124+
append_name=True,
125+
append_path_id=append_path_id,
126+
export_unknown_as_typetree=export_unknown_as_typetree,
127+
)
128+
)
129+
130+
return exported
131+
132+
133+
###############################################################################
134+
# EXPORT FUNCTIONS #
135+
###############################################################################
136+
137+
138+
def exportTextAsset(obj: TextAsset, fp: str, extension: str = ".txt") -> List[int]:
139+
if not extension:
140+
extension = ".txt"
141+
with open(f"{fp}{extension}", "wb") as f:
142+
f.write(obj.script)
143+
return [obj.path_id]
144+
145+
146+
def exportFont(obj: Font, fp: str, extension: str = "") -> List[int]:
147+
# TODO - export glyphs
148+
if obj.m_FontData:
149+
extension = ".ttf"
150+
if obj.m_FontData[0:4] == b"OTTO":
151+
extension = ".otf"
152+
with open(f"{fp}{extension}", "wb") as f:
153+
f.write(obj.m_FontData)
154+
return [obj.path_id]
155+
156+
157+
def exportMesh(obj: Mesh, fp: str, extension=".obf") -> List[int]:
158+
if not extension:
159+
extension = ".obf"
160+
with open(f"{fp}{extension}", "wt", encoding="utf8", newline="") as f:
161+
f.write(obj.export())
162+
return [obj.path_id]
163+
164+
165+
def exporShader(obj: Shader, fp: str, extension=".txt") -> List[int]:
166+
if not extension:
167+
extension = ".txt"
168+
with open(f"{fp}{extension}", "wt", encoding="utf8", newline="") as f:
169+
f.write(obj.export())
170+
return [obj.path_id]
171+
172+
173+
def exportMonoBehaviour(
174+
obj: Union[MonoBehaviour, Object], fp: str, extension: str = ""
175+
) -> List[int]:
176+
# TODO - add generic way to add external typetrees
177+
if obj.serialized_type.nodes:
178+
extension = ".json"
179+
export = json.dumps(obj.read_typetree(), indent=4, ensure_ascii=False).encode(
180+
"utf8", errors="surrogateescape"
181+
)
182+
elif isinstance(obj, MonoBehaviour):
183+
# no set typetree
184+
# check if we have a script
185+
script = obj.m_Script
186+
if script:
187+
# looks like we have a script
188+
script = script.read()
189+
# check if there is a locally stored typetree for it
190+
nodes = MONOBEHAVIOUR_TYPETREES.get(script.m_AssemblyName, {}).get(
191+
script.m_ClassName, None
192+
)
193+
if nodes:
194+
# we have a typetree
195+
# adjust the name
196+
# name = (
197+
# f"{script.m_ClassName}-{obj.name}"
198+
# if obj.name
199+
# else script.m_ClassName
200+
# )
201+
extension = ".json"
202+
export = json.dumps(
203+
obj.read_typetree(nodes), indent=4, ensure_ascii=False
204+
).encode("utf8", errors="surrogateescape")
205+
if not export:
206+
extension = ".bin"
207+
export = obj.raw_data
208+
with open(f"{fp}{extension}", "wb") as f:
209+
f.write(export)
210+
return [obj.path_id]
211+
212+
213+
def exportAudioClip(obj: AudioClip, fp: str, extension: str = "") -> List[int]:
214+
samples = obj.samples
215+
if len(samples) == 0:
216+
pass
217+
elif len(samples) == 1:
218+
with open(f"{fp}.wav", "wb") as f:
219+
f.write(list(samples.values())[0])
220+
else:
221+
os.makedirs(fp, exist_ok=True)
222+
for name, clip_data in samples.items():
223+
with open(os.path.join(fp, f"{name}.wav"), "wb") as f:
224+
f.write(clip_data)
225+
return [obj.path_id]
226+
227+
228+
def exportSprite(obj: Sprite, fp: str, extension: str = ".png") -> List[int]:
229+
if not extension:
230+
extension = ".png"
231+
obj.image.save(f"{fp}{extension}")
232+
return [
233+
obj.path_id,
234+
obj.m_RD.texture.path_id,
235+
getattr(obj.m_RD.alphaTexture, "path_id", None),
236+
]
237+
238+
239+
def exportTexture2D(obj: Texture2D, fp: str, extension: str = ".png") -> List[int]:
240+
if not extension:
241+
extension = ".png"
242+
if obj.m_Width:
243+
# textures can be empty
244+
obj.image.save(f"{fp}{extension}")
245+
return [obj.path_id]
246+
247+
248+
def exportGameObject(obj: GameObject, fp: str, extension: str = "") -> List[int]:
249+
exported = [obj.path_id]
250+
refs = crawl_obj(obj)
251+
if refs:
252+
os.makedirs(fp, exist_ok=True)
253+
for ref_id, ref in refs.items():
254+
# Don't export already exported objects a second time
255+
# and prevent circular calls by excluding other GameObjects.
256+
# The other GameObjects were already exported in the this call.
257+
if ref_id in exported or ref_id.type == ClassIDType.GameObject:
258+
continue
259+
try:
260+
exported.extend(export_obj(ref, fp, True, True))
261+
except Exception as e:
262+
print(f"Failed to export {ref_id}")
263+
print(e)
264+
return exported
265+
266+
267+
EXPORT_TYPES = {
268+
# following types can include other objects
269+
ClassIDType.GameObject: exportGameObject,
270+
ClassIDType.Sprite: exportSprite,
271+
# following types don't include other objects
272+
ClassIDType.AudioClip: exportAudioClip,
273+
ClassIDType.Font: exportFont,
274+
ClassIDType.Mesh: exportMesh,
275+
ClassIDType.MonoBehaviour: exportMonoBehaviour,
276+
ClassIDType.Shader: exporShader,
277+
ClassIDType.TextAsset: exportTextAsset,
278+
ClassIDType.Texture2D: exportTexture2D,
279+
}
280+
281+
MONOBEHAVIOUR_TYPETREES: Dict["Assembly-Name.dll", Dict["Class-Name", List[Dict]]] = {}
282+
283+
284+
def crawl_obj(obj: Object, ret: dict = None) -> Dict[int, Union[Object, PPtr]]:
285+
"""Crawls through the data struture of the object and returns a list of all the components."""
286+
if not ret:
287+
ret = {}
288+
289+
if isinstance(obj, PPtr):
290+
if obj.path_id == 0 and obj.file_id == 0 and obj.index == -2:
291+
return ret
292+
try:
293+
obj = obj.read()
294+
except AttributeError:
295+
return ret
296+
else:
297+
return ret
298+
ret[obj.path_id] = obj
299+
300+
# MonoBehaviour really on their typetree
301+
# while Object denotes that the class of the object isn't implemented yet
302+
if isinstance(obj, (MonoBehaviour, Object)):
303+
obj.read_typetree()
304+
data = obj.type_tree.__dict__.values()
305+
else:
306+
data = obj.__dict__.values()
307+
308+
for value in flatten(data):
309+
if isinstance(value, (Object, PPtr)):
310+
if value.path_id in ret:
311+
continue
312+
crawl_obj(value, ret)
313+
314+
return ret
315+
316+
317+
def flatten(l):
318+
for el in list(l):
319+
if isinstance(el, (list, tuple)):
320+
yield from flatten(el)
321+
elif isinstance(el, dict):
322+
yield from flatten(el.values())
323+
else:
324+
yield el

0 commit comments

Comments
 (0)