from __future__ import annotations

from ntpath import basename

from typing import TYPE_CHECKING, Dict, List, Optional, Tuple

from attrs import define

from .. import config

from ..enums import BuildTarget, ClassIDType

from ..helpers.ContainerHelper import ContainerHelper

from ..helpers.Tpk import get_common_strings

from ..helpers.TypeTreeHelper import TypeTreeNode

from ..helpers.UnityVersion import UnityVersion

from ..streams import EndianBinaryWriter

from . import BundleFile, File

from .ObjectReader import ObjectReader

if TYPE_CHECKING:

from ..classes import AssetBundle, Object

from ..files import ObjectReader

from ..streams.EndianBinaryReader import EndianBinaryReader

@define(slots=True)

class SerializedFileHeader:

metadata_size: int

file_size: int

version: int

data_offset: int

endian: str

reserved: bytes

def __init__(self, reader: EndianBinaryReader):

(

self.metadata_size,

self.file_size,

self.version,

self.data_offset,

) = reader.read_u_int_array(4)

@define(slots=True)

class LocalSerializedObjectIdentifier: # script type

local_serialized_file_index: int

local_identifier_in_file: int

def __init__(self, header: SerializedFileHeader, reader: EndianBinaryReader):

self.local_serialized_file_index = reader.read_int()

if header.version < 14:

self.local_identifier_in_file = reader.read_int()

else:

reader.align_stream()

self.local_identifier_in_file = reader.read_long()

def write(self, header: SerializedFileHeader, writer: EndianBinaryWriter):

writer.write_int(self.local_serialized_file_index)

if header.version < 14:

writer.write_int(self.local_identifier_in_file)

else:

writer.align_stream()

writer.write_long(self.local_identifier_in_file)

@define(slots=True)

class FileIdentifier: # external

path: str

temp_empty: Optional[str] = None

guid: Optional[bytes] = None

type: Optional[int] = None

# enum { kNonAssetType = 0, kDeprecatedCachedAssetType = 1, kSerializedAssetType = 2, kMetaAssetType = 3 };

@property

def name(self):

return basename(self.path)

def __repr__(self):

return f"<{self.__class__.__name__}({self.path})>"

def __init__(self, header: SerializedFileHeader, reader: EndianBinaryReader):

if header.version >= 6:

self.temp_empty = reader.read_string_to_null()

if header.version >= 5:

self.guid = reader.read_bytes(16)

self.type = reader.read_int()

self.path = reader.read_string_to_null()

def write(self, header: SerializedFileHeader, writer: EndianBinaryWriter):

if header.version >= 6:

assert self.temp_empty is not None

writer.write_string_to_null(self.temp_empty)

if header.version >= 5:

assert self.guid is not None and self.type is not None

writer.write_bytes(self.guid)

writer.write_int(self.type)

writer.write_string_to_null(self.path)

@define(slots=True, init=False)

class SerializedType:

class_id: int

is_stripped_type: Optional[bool] = None

script_type_index: int = -1

script_id: Optional[bytes] = None # Hash128

old_type_hash: Optional[bytes] = None # Hash128

node: Optional[TypeTreeNode] = None

# ref type

m_ClassName: Optional[str] = None

m_NameSpace: Optional[str] = None

m_AssemblyName: Optional[str] = None

# 21+

type_dependencies: Optional[Tuple[int, ...]] = None

def __init__(

self,

reader: EndianBinaryReader,

serialized_file: SerializedFile,

is_ref_type: bool,

):

version = serialized_file.header.version

self.class_id = reader.read_int()

self.__attrs_init__(self.class_id)

if version >= 16:

self.is_stripped_type = reader.read_boolean()

if version >= 17:

self.script_type_index = reader.read_short()

if version >= 13:

if (

(is_ref_type and self.script_type_index >= 0)

or (version < 16 and self.class_id < 0)

or (version >= 16 and self.class_id == 114)

):

self.script_id = reader.read_bytes(16)

self.old_type_hash = reader.read_bytes(16)

if serialized_file._enable_type_tree:

if version >= 12 or version == 10:

self.node = TypeTreeNode.parse_blob(reader, version)

else:

self.node = TypeTreeNode.parse(reader, version)

if version >= 21:

if is_ref_type:

self.m_ClassName = reader.read_string_to_null()

self.m_NameSpace = reader.read_string_to_null()

self.m_AssemblyName = reader.read_string_to_null()

else:

self.type_dependencies = reader.read_int_array()

def write(

self,

serialized_file: SerializedFile,

writer: EndianBinaryWriter,

is_ref_type: bool,

):

version = serialized_file.header.version

writer.write_int(self.class_id)

if version >= 16:

assert self.is_stripped_type is not None

writer.write_boolean(self.is_stripped_type)

if version >= 17:

assert self.script_type_index is not None

writer.write_short(self.script_type_index)

if version >= 13:

if (

(is_ref_type and self.script_type_index >= 0)

or (version < 16 and self.class_id < 0)

or (version >= 16 and self.class_id == 114)

):

assert self.script_id is not None

writer.write_bytes(self.script_id) # Hash128

assert self.old_type_hash is not None

writer.write_bytes(self.old_type_hash) # Hash128

if serialized_file._enable_type_tree:

assert self.node is not None

if version >= 12 or version == 10:

self.node.dump_blob(writer, version)

else:

serialized_file.dump(writer, version)

if version >= 21:

if is_ref_type:

assert (

self.m_ClassName is not None

and self.m_NameSpace is not None

and self.m_AssemblyName is not None

)

writer.write_string_to_null(self.m_ClassName)

writer.write_string_to_null(self.m_NameSpace)

writer.write_string_to_null(self.m_AssemblyName)

else:

assert self.type_dependencies is not None

writer.write_int_array(self.type_dependencies, True)

@property

def nodes(self) -> Optional[TypeTreeNode]:

# for compatibility with old versions

return self.node

class SerializedFile(File.File):

reader: EndianBinaryReader

version: UnityVersion

unity_version: str

target_platform: BuildTarget

_enable_type_tree: bool

types: List[SerializedType]

script_types: List[LocalSerializedObjectIdentifier]

externals: List[FileIdentifier]

ref_types: Optional[List[SerializedType]]

objects: Dict[int, ObjectReader]

unknown: int

header: SerializedFileHeader

_m_target_platform: int

big_id_enabled: int

userInformation: Optional[str]

assetbundle: Optional[AssetBundle]

_cache: Dict[str, Object]

@property

def files(self):

if self.objects:

return self.objects

return {}

@files.setter

def files(self, value):

self.objects = value

def __init__(self, reader: EndianBinaryReader, parent=None, name=None, **kwargs):

super().__init__(parent=parent, name=name, **kwargs)

self.reader = reader

self.unity_version = "2.5.0f5"

self.target_platform = BuildTarget.UnknownPlatform

self._enable_type_tree = True

self.types = []

self.script_types = []

self.externals = []

self.objects = {}

# used to speed up mass asset extraction

# some assets refer to each other, so by keeping the result

# of specific assets cached the extraction can be speed up by a lot.

# used by: Sprite (Texture2D (with alpha) cached),

self._cache = {}

self.unknown = 0

# ReadHeader

header = SerializedFileHeader(reader)

self.header = header

if header.version >= 9:

header.endian = ">" if reader.read_boolean() else "<"

header.reserved = reader.read_bytes(3)

if header.version >= 22:

header.metadata_size = reader.read_u_int()

header.file_size = reader.read_long()

header.data_offset = reader.read_long()

self.unknown = reader.read_long() # unknown

else:

reader.Position = header.file_size - header.metadata_size

header.endian = ">" if reader.read_boolean() else "<"

reader.endian = header.endian

if header.version >= 7:

unity_version = reader.read_string_to_null()

self.set_version(unity_version)

if header.version >= 8:

self._m_target_platform = reader.read_int()

self.target_platform = BuildTarget(self._m_target_platform)

if header.version >= 13:

self._enable_type_tree = reader.read_boolean()

# ReadTypes

type_count = reader.read_int()

self.types = [SerializedType(reader, self, False) for _ in range(type_count)]

self.big_id_enabled = 0

if 7 <= header.version < 14:

self.big_id_enabled = reader.read_int()

# ReadObjects

object_count = reader.read_int()

self.objects = {}

for _ in range(object_count):

obj = ObjectReader.from_reader(self, reader)

self.objects[obj.path_id] = obj

# Read Scripts

if header.version >= 11:

script_count = reader.read_int()

self.script_types = [LocalSerializedObjectIdentifier(header, reader) for _ in range(script_count)]

# Read Externals

externals_count = reader.read_int()

self.externals = [FileIdentifier(header, reader) for _ in range(externals_count)]

if header.version >= 20:

ref_type_count = reader.read_int()

self.ref_types = [SerializedType(reader, self, True) for _ in range(ref_type_count)]

if config.SERIALIZED_FILE_PARSE_TYPETREE is False:

self._enable_type_tree = False

if header.version >= 5:

self.userInformation = reader.read_string_to_null()

# read the asset_bundles to get the containers

for obj in self.objects.values():

if obj.type == ClassIDType.AssetBundle:

self.assetbundle = obj.parse_as_object()

self._container = ContainerHelper(self.assetbundle)

break

else:

self.assetbundle = None

self._container = ContainerHelper([])

@property

def container(self):

return self._container

def load_dependencies(self, possible_dependencies: Optional[list] = None):

"""Load all external dependencies.

Parameters

----------

possible_dependencies : list

List of possible dependencies for cases

where the target file is not listed as external.

"""

for file_id in self.externals:

self.environment.load_file(file_id.path, True)

if possible_dependencies is None:

return

for dependency in possible_dependencies:

try:

self.environment.load_file(dependency, True)

except FileNotFoundError:

pass

def set_version(self, string_version: str):

self.unity_version = string_version

if not string_version or string_version == "0.0.0":

# weird case, but apparently can happen?

# check "cant read Texture2D by 2020.3.13 f1 AssetBundle #77" for details

if isinstance(self.parent, BundleFile.BundleFile):

string_version = self.parent.version_engine

if not string_version or string_version == "0.0.0":

string_version = config.get_fallback_version()

self.version = UnityVersion.from_str(string_version)

def get_writeable_cab(self, name: str = "CAB-UnityPy_Mod.resS"):

"""

Creates a new cab file in the bundle that contains the given data.

This is usefull for asset types that use resource files.

"""

if not isinstance(self.parent, (File.BundleFile.BundleFile, File.WebFile.WebFile)):

return None

cab = self.parent.get_writeable_cab(name)

cab.path = f"archive:/{self.name}/{name}"

if not any(cab.path == x.path for x in self.externals):

# register as external

class FileIdentifierFake:

pass

file_identifier = FileIdentifierFake()

file_identifier.__class__ = FileIdentifier

file_identifier.temp_empty = ""

import uuid

file_identifier.guid = uuid.uuid1().urn[-16:].encode("ascii")

file_identifier.path = cab.path

file_identifier.type = 0

self.externals.append(file_identifier)

return cab

def save(self, packer: Optional[str] = None) -> bytes:

# 1. header -> has to be delayed until the very end

# 2. data -> types, objects, scripts, ...

# so write the data first

header = self.header

meta_writer = EndianBinaryWriter(endian=header.endian)

data_writer = EndianBinaryWriter(endian=header.endian)

if header.version >= 7:

meta_writer.write_string_to_null(self.unity_version)

if header.version >= 8:

meta_writer.write_int(self._m_target_platform)

if header.version >= 13:

meta_writer.write_boolean(self._enable_type_tree)

# ReadTypes

meta_writer.write_int(len(self.types))

for typ in self.types:

typ.write(self, meta_writer, False)

if 7 <= header.version < 14:

meta_writer.write_int(self.big_id_enabled)

# ReadObjects

meta_writer.write_int(len(self.objects))

for obj in sorted(self.objects.values(), key=lambda x: x.path_id):

obj.write(header, meta_writer, data_writer)

data_writer.align_stream(8)

# Read Scripts

if header.version >= 11:

meta_writer.write_int(len(self.script_types))

for script_type in self.script_types:

script_type.write(header, meta_writer)

# Read Externals

meta_writer.write_int(len(self.externals))

for external in self.externals:

external.write(header, meta_writer)

if header.version >= 20:

assert self.ref_types is not None

meta_writer.write_int(len(self.ref_types))

for ref_type in self.ref_types:

ref_type.write(self, meta_writer, True)

if header.version >= 5:

assert self.userInformation is not None

meta_writer.write_string_to_null(self.userInformation)

# prepare header

writer = EndianBinaryWriter()

header_size = 16 # 4*4

metadata_size = meta_writer.Length

data_size = data_writer.Length

if header.version >= 9:

# 1 bool + 3 reserved + extra header 4 + 3*8

header_size += 4 if header.version < 22 else 4 + 28

data_offset = header_size + metadata_size

# align data_offset

data_offset += (16 - data_offset % 16) % 16

file_size = data_offset + data_size

if header.version < 22:

writer.write_u_int(metadata_size)

writer.write_u_int(file_size)

writer.write_u_int(header.version)

# reader.Position = header.file_size - header.metadata_size

# so data follows right after this header -> after 32

writer.write_u_int(data_offset)

writer.write_boolean(">" == header.endian)

writer.write_bytes(header.reserved)

else:

# old header

writer.write_u_int(0)

writer.write_u_int(0)

writer.write_u_int(header.version)

writer.write_u_int(0)

writer.write_boolean(">" == header.endian)

writer.write_bytes(header.reserved)

writer.write_u_int(metadata_size)

writer.write_long(file_size)

writer.write_long(data_offset)

writer.write_long(self.unknown)

writer.write_bytes(meta_writer.bytes)

writer.align_stream(16)

writer.write_bytes(data_writer.bytes)

else:

metadata_size += 1 # endian boolean

file_size = header_size + metadata_size + data_size

writer.write_u_int(metadata_size)

writer.write_u_int(file_size)

writer.write_u_int(header.version)

# reader.Position = header.file_size - header.metadata_size

# so data follows right after this header -> after 32

writer.write_u_int(32)

writer.write_bytes(data_writer.bytes)

writer.write_boolean(">" == header.endian)

writer.write_bytes(meta_writer.bytes)

return writer.bytes

def read_string(string_buffer_reader: EndianBinaryReader, value: int) -> str:

is_offset = (value & 0x80000000) == 0

if is_offset:

string_buffer_reader.Position = value

return string_buffer_reader.read_string_to_null()

offset = value & 0x7FFFFFFF

return get_common_strings().get(offset, str(offset))