import json import os from io import BytesIO from pathlib import Path from typing import Callable, Dict, Iterable, List, Optional, Sequence, Tuple, Union import UnityPy from UnityPy.classes import ( AudioClip, Font, GameObject, Mesh, MonoBehaviour, Object, PPtr, Shader, Sprite, TextAsset, Texture2D, ) from UnityPy.enums.ClassIDType import ClassIDType from UnityPy.files import SerializedFile def export_obj( obj: Union[Object, PPtr], fp: Path, append_name: bool = False, append_path_id: bool = False, export_unknown_as_typetree: bool = False, asset_filter: Optional[Callable[[Object], bool]] = None, ) -> List[Tuple[SerializedFile, int]]: """Exports the given object to the given filepath. Args: obj (Object, PPtr): A valid Unity object or a reference to one. fp (Path): A valid filepath where the object should be exported to. append_name (bool, optional): Decides if the obj name will be appended to the filepath. Defaults to False. append_path_id (bool, optional): Decides if the obj path id will be appended to the filepath. Defaults to False. export_unknown_as_typetree (bool, optional): If set, then unimplemented objects will be exported via their typetree or dumped as bin. Defaults to False. asset_filter (func(Object)->bool, optional): Determines whether to export an object. Defaults to all objects. Returns: list: a list of exported object path_ids """ # figure out export function export_func = EXPORT_TYPES.get(obj.type) if not export_func: if export_unknown_as_typetree: export_func = exportMonoBehaviour else: return [] # set filepath obj = obj.read() # a filter that returned True during an earlier extract_assets check can return False now with more info from read() if asset_filter and not asset_filter(obj): return [] if append_name: fp = os.path.join( fp, obj.m_Name if getattr(obj, "m_Name") else obj.object_reader.type.name, # noqa: B009 ) fp, extension = os.path.splitext(fp) if append_path_id: fp = f"{fp}_{obj.object_reader.path_id}" # export return export_func(obj, fp, extension) def extract_assets( src: Union[Path, BytesIO, bytes, bytearray], dst: Path, use_container: bool = True, ignore_first_container_dirs: int = 0, append_path_id: bool = False, export_unknown_as_typetree: bool = False, asset_filter: Optional[Callable[[Object], bool]] = None, ) -> List[Tuple[SerializedFile, int]]: """Extracts some or all assets from the given source. Args: src (Union[Path, BytesIO, bytes, bytearray]): [description] dst (Path): [description] use_container (bool, optional): [description]. Defaults to True. ignore_first_container_dirs (int, optional): [description]. Defaults to 0. append_path_id (bool, optional): [description]. Defaults to False. export_unknown_as_typetree (bool, optional): [description]. Defaults to False. asset_filter (func(object)->bool, optional): Determines whether to export an object. Defaults to all objects. Returns: List[Tuple[SerializedFile, int]]: [description] """ # load source env = UnityPy.load(src) exported = [] export_types_keys = list(EXPORT_TYPES.keys()) def defaulted_export_index(type: ClassIDType): try: return export_types_keys.index(type) except (IndexError, ValueError): return 999 if use_container: container = sorted(env.container.items(), key=lambda x: defaulted_export_index(x[1].type)) for obj_path, obj in container: # The filter here can only access metadata. # The same filter may produce a different result later in extract_obj after obj.read() if asset_filter is not None and not asset_filter(obj): continue # the check of the various sub directories is required to avoid // in the path obj_dest = os.path.join( dst, *(x for x in obj_path.split("/")[ignore_first_container_dirs:] if x), ) os.makedirs(os.path.dirname(obj_dest), exist_ok=True) exported.extend( export_obj( obj, obj_dest, append_path_id=append_path_id, export_unknown_as_typetree=export_unknown_as_typetree, asset_filter=asset_filter, ) ) else: objects = sorted(env.objects, key=lambda x: defaulted_export_index(x.type)) for obj in objects: if asset_filter is not None and not asset_filter(obj): continue if (obj.assets_file, obj.path_id) not in exported: exported.extend( export_obj( obj, dst, append_name=True, append_path_id=append_path_id, export_unknown_as_typetree=export_unknown_as_typetree, asset_filter=asset_filter, ) ) return exported ############################################################################### # EXPORT FUNCTIONS # ############################################################################### def exportTextAsset(obj: TextAsset, fp: str, extension: str = ".txt") -> List[Tuple[SerializedFile, int]]: if not extension: extension = ".txt" with open(f"{fp}{extension}", "wb") as f: f.write(obj.m_Script.encode("utf-8", "surrogateescape")) return [(obj.assets_file, obj.object_reader.path_id)] def exportFont(obj: Font, fp: str, extension: str = "") -> List[Tuple[SerializedFile, int]]: # TODO - export glyphs if obj.m_FontData: extension = ".ttf" if obj.m_FontData[0:4] == b"OTTO": extension = ".otf" with open(f"{fp}{extension}", "wb") as f: f.write(bytes(obj.m_FontData)) return [(obj.assets_file, obj.object_reader.path_id)] def exportMesh(obj: Mesh, fp: str, extension=".obj") -> List[Tuple[SerializedFile, int]]: if not extension: extension = ".obj" with open(f"{fp}{extension}", "wt", encoding="utf8", newline="") as f: f.write(obj.export()) return [(obj.assets_file, obj.object_reader.path_id)] def exportShader(obj: Shader, fp: str, extension=".txt") -> List[Tuple[SerializedFile, int]]: if not extension: extension = ".txt" with open(f"{fp}{extension}", "wt", encoding="utf8", newline="") as f: f.write(obj.export()) return [(obj.assets_file, obj.object_reader.path_id)] def exportMonoBehaviour( obj: Union[MonoBehaviour, Object], fp: str, extension: str = "" ) -> List[Tuple[SerializedFile, int]]: export = None if obj.object_reader.serialized_type.node: # a typetree is available from the SerializedFile for this object export = obj.object_reader.read_typetree() elif isinstance(obj, MonoBehaviour): # try to get the typetree from the MonoBehavior script script_ptr = obj.m_Script if script_ptr: # looks like we have a script script = script_ptr.read() # check if there is a locally stored typetree for it nodes = MONOBEHAVIOUR_TYPETREES.get(script.m_AssemblyName, {}).get(script.m_ClassName, None) if nodes: export = obj.object_reader.read_typetree(nodes) else: export = obj.object_reader.read_typetree() if not export: extension = ".bin" export = obj.object_reader.raw_data else: extension = ".json" export = json.dumps(export, indent=4, ensure_ascii=False).encode("utf8", errors="surrogateescape") with open(f"{fp}{extension}", "wb") as f: f.write(export) return [(obj.assets_file, obj.object_reader.path_id)] def exportAudioClip(obj: AudioClip, fp: str, extension: str = "") -> List[Tuple[SerializedFile, int]]: samples = obj.samples if len(samples) == 0: pass elif len(samples) == 1: with open(f"{fp}.wav", "wb") as f: f.write(list(samples.values())[0]) else: os.makedirs(fp, exist_ok=True) for name, clip_data in samples.items(): with open(os.path.join(fp, f"{name}.wav"), "wb") as f: f.write(clip_data) return [(obj.assets_file, obj.object_reader.path_id)] def exportSprite(obj: Sprite, fp: str, extension: str = ".png") -> List[Tuple[SerializedFile, int]]: if not extension: extension = ".png" obj.image.save(f"{fp}{extension}") exported = [ (obj.assets_file, obj.object_reader.path_id), (obj.m_RD.texture.assetsfile, obj.m_RD.texture.path_id), ] alpha_assets_file = getattr(obj.m_RD.alphaTexture, "assets_file", None) alpha_path_id = getattr(obj.m_RD.alphaTexture, "path_id", None) if alpha_path_id and alpha_assets_file: exported.append((alpha_assets_file, alpha_path_id)) return exported def exportTexture2D(obj: Texture2D, fp: str, extension: str = ".png") -> List[Tuple[SerializedFile, int]]: if not extension: extension = ".png" if obj.m_Width: # textures can be empty obj.image.save(f"{fp}{extension}") return [(obj.assets_file, obj.object_reader.path_id)] def exportGameObject(obj: GameObject, fp: str, extension: str = "") -> List[Tuple[SerializedFile, int]]: exported = [(obj.assets_file, obj.object_reader.path_id)] refs = crawl_obj(obj) if refs: os.makedirs(fp, exist_ok=True) for ref_id, ref in refs.items(): # Don't export already exported objects a second time # and prevent circular calls by excluding other GameObjects. # The other GameObjects were already exported in the this call. if (ref.assets_file, ref_id) in exported or ref.type == ClassIDType.GameObject: continue try: exported.extend(export_obj(ref, fp, True, True)) except Exception as e: print(f"Failed to export {ref_id}") print(e) return exported EXPORT_TYPES = { # following types can include other objects ClassIDType.GameObject: exportGameObject, ClassIDType.Sprite: exportSprite, # following types don't include other objects ClassIDType.AudioClip: exportAudioClip, ClassIDType.Font: exportFont, ClassIDType.Mesh: exportMesh, ClassIDType.MonoBehaviour: exportMonoBehaviour, ClassIDType.Shader: exportShader, ClassIDType.TextAsset: exportTextAsset, ClassIDType.Texture2D: exportTexture2D, } ASSEMBLY_NAME_DLL = str CLASS_NAME = str MONOBEHAVIOUR_TYPETREES: Dict[ASSEMBLY_NAME_DLL, Dict[CLASS_NAME, List[Dict]]] = {} def crawl_obj(obj: Object, ret: Optional[dict] = None) -> Dict[int, Union[Object, PPtr]]: """Crawls through the data struture of the object and returns a list of all the components. """ if not ret: ret = {} if isinstance(obj, PPtr): if obj.path_id == 0 and obj.file_id == 0 and obj.index == -2: return ret try: obj = obj.read() except AttributeError: return ret else: return ret ret[obj.path_id] = obj # MonoBehaviour really on their typetree # while Object denotes that the class of the object isn't implemented yet if isinstance(obj, (MonoBehaviour, Object)): data = obj.read_typetree().__dict__.values() else: data = obj.__dict__.values() for value in flatten(data): if isinstance(value, (Object, PPtr)): if value.path_id in ret: continue crawl_obj(value, ret) return ret def flatten(seq: Sequence) -> Iterable: for elem in list(seq): if isinstance(elem, (list, tuple)): yield from flatten(elem) elif isinstance(elem, dict): yield from flatten(elem.values()) # type: ignore else: yield elem