UnityPy/UnityPy/helpers/Tpk.py at docs · isHarryh/UnityPy

473 lines (375 loc) · 13.9 KB
from __future__ import annotations
from enum import IntEnum, IntFlag
from importlib.resources import open_binary
from io import BytesIO
from struct import Struct
from typing import Any, Dict, List, Optional, Tuple, TypeVar
from .CompressionHelper import decompress_lzma
from .TypeTreeHelper import TypeTreeNode
from .UnityVersion import UnityVersion
T = TypeVar("T")
TPKTYPETREE: TpkTypeTreeBlob = None  # pyright: ignore[reportAssignmentType]
CLASSES_CACHE: Dict[Tuple[int, UnityVersion], TypeTreeNode] = {}
NODES_CACHE: Dict[TpkUnityClass, TypeTreeNode] = {}
def init():
    with open_binary("UnityPy.resources", "lzma.tpk") as f:
        data = f.read()
    global TPKTYPETREE
    with BytesIO(data) as stream:
        blob = TpkFile(stream).GetDataBlob()
        assert isinstance(blob, TpkTypeTreeBlob)
        TPKTYPETREE = blob
def get_typetree_node(class_id: int, version: UnityVersion):
    global CLASSES_CACHE
    key = (class_id, version)
    cached = CLASSES_CACHE.get(key)
    if cached:
        return cached
    class_info = TPKTYPETREE.ClassInformation[class_id].getVersionedClass(version)
    if class_info is None:
        raise ValueError("Could not find class info for class id {}".format(class_id))
    node = generate_node(class_info)
    CLASSES_CACHE[key] = node
    return node
def generate_node(class_info: TpkUnityClass) -> TypeTreeNode:
    global NODES_CACHE
    cached = NODES_CACHE.get(class_info)
    if cached:
        return cached
    assert class_info.ReleaseRootNode is not None, "Class {} has no ReleaseRootNode".format(class_info)
    nodes = []
    NODES = TPKTYPETREE.NodeBuffer
    stack = [(class_info.ReleaseRootNode, 0)]
    index = 0
    while stack:
        node_id, level = stack.pop(0)
        node: TpkUnityNode = NODES[node_id]
        nodes.append(
            TypeTreeNode(
                m_ByteSize=node.ByteSize,
                m_Index=index,
                m_Version=node.Version,
                m_MetaFlag=node.MetaFlag,
                m_Level=level,
                m_Type=TPKTYPETREE.StringBuffer[node.TypeName],
                m_Name=TPKTYPETREE.StringBuffer[node.Name],
        stack = [(node_id, level + 1) for node_id in node.SubNodes] + stack
        index += 1
    result = TypeTreeNode.from_list(nodes)
    NODES_CACHE[class_info] = result
    return result
######################################################################################
######################################################################################
class TpkCompressionType(IntEnum):
    NONE = 0
    Lz4 = 1
    Lzma = 2
    Brotli = 3
class TpkDataType(IntEnum):
    TypeTreeInformation = 0
    Collection = 1
    FileSystem = 2
    Json = 3
    ReferenceAssemblies = 4
    EngineAssets = 5
    def ToBlob(self, stream):
        if self.value == TpkDataType.TypeTreeInformation:
            return TpkTypeTreeBlob(stream)
        elif self.value == TpkDataType.Collection:
            return TpkCollectionBlob(stream)
        elif self.value == TpkDataType.FileSystem:
            return TpkFileSystemBlob(stream)
        elif self.value == TpkDataType.Json:
            return TpkJsonBlob(stream)
        else:
            raise Exception("Unimplemented TpkDataType -> Blob conversion")
class TpkUnityClassFlags(IntFlag):
    NONE = 0
    IsAbstract = 1
    IsSealed = 2
    IsEditorOnly = 4
    IsReleaseOnly = 8
    IsStripped = 16
    Reserved = 32
    HasEditorRootNode = 64
    HasReleaseRootNode = 128
######################################################################################
#   Main Class
######################################################################################
class TpkFile:
    Struct = Struct("<IbbbbIII")
    TpkMagicBytes: int = 0x2A4B5054  # b"TPK*"
    TpkVersionNumber: int = 1
    CompressionType: TpkCompressionType
    DataType: TpkDataType
    CompressedSize: int
    UncompressedSize: int
    CompressedBytes: bytes
    def __init__(self, stream: BytesIO):
            magic,
            versionNumber,
            compressionType,
            dataType,
            self.CompressedSize,
            self.UncompressedSize,
        ) = TpkFile.Struct.unpack(stream.read(TpkFile.Struct.size))
        if magic != TpkFile.TpkMagicBytes:
            raise Exception("Invalid TPK magic bytes")
        if versionNumber != TpkFile.TpkVersionNumber:
            raise Exception("Invalid TPK version number")
        self.CompressionType = TpkCompressionType(compressionType)
        self.DataType = TpkDataType(dataType)
        self.CompressedBytes = stream.read(self.CompressedSize)
        if len(self.CompressedBytes) != self.CompressedSize:
            raise Exception("Invalid compressed size")
    def GetDataBlob(self) -> TpkDataBlob:
        decompressed: bytes
        if self.CompressionType == TpkCompressionType.NONE:
            decompressed = self.CompressedBytes
        elif self.CompressionType == TpkCompressionType.Lz4:
            import lz4.block
            decompressed = lz4.block.decompress(self.CompressedBytes, self.UncompressedSize)
        elif self.CompressionType == TpkCompressionType.Lzma:
            decompressed = decompress_lzma(self.CompressedBytes)
        elif self.CompressionType == TpkCompressionType.Brotli:
            import brotli
            decompressed = brotli.decompress(self.CompressedBytes)
        else:
            raise Exception("Invalid compression type")
        return self.DataType.ToBlob(BytesIO(decompressed))
######################################################################################
######################################################################################
class TpkDataBlob:
    __slots__ = ("DataType",)
    DataType: TpkDataType
    def __init__(self, stream: BytesIO) -> None:
        raise NotImplementedError("TpkDataBlob is an abstract class")
class TpkTypeTreeBlob(TpkDataBlob):
    __slots__ = (
        "CreationTime",
        "Versions",
        "ClassInformation",
        "CommonString",
        "NodeBuffer",
        "StringBuffer",
    CreationTime: int
    Versions: List[UnityVersion]
    ClassInformation: Dict[int, TpkClassInformation]  # List[TpkClassInformation]
    CommonString: TpkCommonString
    NodeBuffer: TpkUnityNodeBuffer
    StringBuffer: TpkStringBuffer
    DataType: TpkDataType = TpkDataType.TypeTreeInformation
    def __init__(self, stream: BytesIO) -> None:
        (self.CreationTime,) = INT64.unpack(stream.read(INT64.size))
        (versionCount,) = INT32.unpack(stream.read(INT32.size))
        self.Versions = read_versions(stream, versionCount)
        (classCount,) = INT32.unpack(stream.read(INT32.size))
        self.ClassInformation = {x.ID: x for x in (TpkClassInformation(stream) for _ in range(classCount))}
        self.CommonString = TpkCommonString(stream)
        self.NodeBuffer = TpkUnityNodeBuffer(stream)
        self.StringBuffer = TpkStringBuffer(stream)
class TpkCollectionBlob(TpkDataBlob):
    __slots__ = "Blobs"
    Blobs: List[Tuple[str, TpkDataBlob]]
    def __init__(self, stream: BytesIO) -> None:
        (count,) = INT32.unpack(stream.read(INT32.size))
        self.Blobs = [
            # relativePath, data
                read_string(stream),
                TpkDataType(BYTE.unpack(stream.read(1))[0]).ToBlob(stream),
            for _ in range(count)
class TpkFileSystemBlob(TpkDataBlob):
    __slots__ = ("Files",)
    # TODO: check if dict might be better
    Files: List[Tuple[str, bytes]]
    def __init__(self, stream: BytesIO) -> None:
        (count,) = INT32.unpack(stream.read(INT32.size))
        self.Files = [
            # relativePath, data
            (read_string(stream), read_data(stream))
            for _ in range(count)
class TpkJsonBlob(TpkDataBlob):
    __slots__ = "Text"
    Text: str
    DataType = TpkDataType.Json
    def __init__(self, stream: BytesIO) -> None:
        self.Text = read_string(stream)
######################################################################################
######################################################################################
class TpkUnityClass:
    __slots__ = ("Name", "Base", "Flags", "EditorRootNode", "ReleaseRootNode")
    Struct = Struct("<HHb")
    Name: int
    Base: int
    Flags: TpkUnityClassFlags
    EditorRootNode: Optional[int]
    ReleaseRootNode: Optional[int]
    def __init__(self, stream: BytesIO) -> None:
        self.Name, self.Base, Flags = TpkUnityClass.Struct.unpack(stream.read(TpkUnityClass.Struct.size))
        self.Flags = TpkUnityClassFlags(Flags)
        self.EditorRootNode = self.ReleaseRootNode = None
        if self.Flags & TpkUnityClassFlags.HasEditorRootNode:
            (self.EditorRootNode,) = UINT16.unpack(stream.read(UINT16.size))
        if self.Flags & TpkUnityClassFlags.HasReleaseRootNode:
            (self.ReleaseRootNode,) = UINT16.unpack(stream.read(UINT16.size))
    def to_dict(self) -> Dict[str, Any]:
        return {
            "Name": self.Name,
            "Base": self.Base,
            "Flags": self.Flags,
            "EditorRootNode": self.EditorRootNode,
            "ReleaseRootNode": self.ReleaseRootNode,
    def __eq__(self, other: object) -> bool:
        if not isinstance(other, TpkUnityClass):
            return False
        return self.to_dict() == other.to_dict()
    def __hash__(self) -> int:
        return hash(
                self.Name,
                self.Base,
                self.Flags,
                self.EditorRootNode,
                self.ReleaseRootNode,
class TpkClassInformation(List[Tuple[UnityVersion, Optional[TpkUnityClass]]]):
    ID: int
    def __init__(self, stream: BytesIO) -> None:
        (self.ID,) = INT32.unpack(stream.read(INT32.size))
        (count,) = INT32.unpack(stream.read(INT32.size))
        self.extend(
                read_version(stream),
                TpkUnityClass(stream) if stream.read(1)[0] else None,
            for _ in range(count)
    def getVersionedClass(self, version: UnityVersion) -> Optional[TpkUnityClass]:
        return get_item_for_version(version, self)
class TpkUnityNode:
    __slots__ = (
        "TypeName",
        "Name",
        "ByteSize",
        "Version",
        "TypeFlags",
        "MetaFlag",
        "SubNodes",
    Struct = Struct("<HHihbIH")
    TypeName: int
    Name: int
    ByteSize: int
    Version: int
    TypeFlags: int
    MetaFlag: int
    SubNodes: List[int]
    def __init__(self, stream: BytesIO) -> None:
            self.TypeName,
            self.Name,
            self.ByteSize,
            self.Version,
            self.TypeFlags,
            self.MetaFlag,
            count,
        ) = TpkUnityNode.Struct.unpack(stream.read(TpkUnityNode.Struct.size))
        SubNodeStruct = Struct(f"<{count}H")
        self.SubNodes = list(SubNodeStruct.unpack(stream.read(SubNodeStruct.size)))
    def __eq__(self, other: object) -> bool:
        if not isinstance(other, TpkUnityNode):
            return False
        return self.__dict__ == other.__dict__
    def __hash__(self) -> int:
        # TODO
        return hash(self.__dict__)
class TpkUnityNodeBuffer(List[TpkUnityNode]):
    def __init__(self, stream: BytesIO) -> None:
        (count,) = INT32.unpack(stream.read(INT32.size))
        self.extend(TpkUnityNode(stream) for _ in range(count))
######################################################################################
#   Strings
######################################################################################
class TpkStringBuffer(List[str]):
    def __init__(self, stream: BytesIO) -> None:
        count = INT32.unpack(stream.read(INT32.size))[0]
        self.extend(read_string(stream) for _ in range(count))
class TpkCommonString:
    __slots__ = ("VersionInformation", "StringBufferIndices")
    VersionInformation: List[Tuple[UnityVersion, int]]
    StringBufferIndices: Tuple[int, ...]
    def __init__(self, stream: BytesIO) -> None:
        (versionCount,) = INT32.unpack(stream.read(INT32.size))
        self.VersionInformation = [(read_version(stream), stream.read(1)[0]) for _ in range(versionCount)]
        (indicesCount,) = INT32.unpack(stream.read(INT32.size))
        indicesStruct = Struct(f"<{indicesCount}H")
        self.StringBufferIndices = indicesStruct.unpack(stream.read(indicesStruct.size))
    def GetStrings(self, buffer: TpkStringBuffer) -> List[str]:
        return [buffer[i] for i in self.StringBufferIndices]
    def GetCount(self, exactVersion: UnityVersion) -> int:
        return get_item_for_version(exactVersion, self.VersionInformation)
######################################################################################
# helper functions
######################################################################################
BYTE = Struct("b")
UINT16 = Struct("<H")
INT32 = Struct("<i")
INT64 = Struct("<q")
UINT64 = Struct("<Q")
def read_string(stream: BytesIO) -> str:
    # varint
    shift = 0
    length = 0
    while True:
        (i,) = stream.read(1)
        length |= (i & 0x7F) << shift
        shift += 7
        if not (i & 0x80):
            break
    # string
    return stream.read(length).decode("utf-8")
def read_data(stream: BytesIO) -> bytes:
    return stream.read(INT32.unpack(stream.read(INT32.size))[0])
def read_version(stream: BytesIO) -> UnityVersion:
    return UnityVersion(UINT64.unpack(stream.read(UINT64.size))[0])
def read_versions(stream: BytesIO, count: int) -> List[UnityVersion]:
    struct = Struct(f"<{count}Q")
    return [UnityVersion(x) for x in struct.unpack(stream.read(struct.size))]
def get_item_for_version(exactVersion: UnityVersion, items: List[Tuple[UnityVersion, T]]) -> T:
    ret = None
    for version, item in items:
        if exactVersion >= version:
            ret = item
        else:
            break
    if ret:
        return ret
    raise ValueError("Could not find exact version")
Provide feedback

Saved searches

Use saved searches to filter your results more quickly

FilesExpand file tree

Tpk.py

Latest commit

History

Tpk.py

File metadata and controls
