diff --git a/UnityPy/helpers/ArchiveStorageManager.py b/UnityPy/helpers/ArchiveStorageManager.py index 797aa018..3b8e3f65 100644 --- a/UnityPy/helpers/ArchiveStorageManager.py +++ b/UnityPy/helpers/ArchiveStorageManager.py @@ -1,6 +1,6 @@ # based on: https://github.com/Razmoth/PGRStudio/blob/master/AssetStudio/PGR/PGR.cs import re -from typing import Tuple, Union +from typing import Optional, Tuple, Union from ..streams import EndianBinaryReader @@ -10,7 +10,7 @@ UnityPyBoost = None UNITY3D_SIGNATURE = b"#$unity3dchina!@" -DECRYPT_KEY: bytes = None +DECRYPT_KEY: Optional[bytes] = None def set_assetbundle_decrypt_key(key: Union[bytes, str]): @@ -61,20 +61,12 @@ def brute_force_key( return None -def to_uint4_array(source: bytes, offset: int = 0): - buffer = bytearray(len(source) * 2) - for j in range(len(source)): - buffer[j * 2] = source[offset + j] >> 4 - buffer[j * 2 + 1] = source[offset + j] & 15 - return buffer - - class ArchiveStorageDecryptor: unknown_1: int index: bytes substitute: bytes = bytes(0x10) - def __init__(self, reader: EndianBinaryReader) -> None: + def __init__(self, reader: EndianBinaryReader): self.unknown_1 = reader.read_u_int() # read vector data/key vectors @@ -99,14 +91,15 @@ def __init__(self, reader: EndianBinaryReader) -> None: raise Exception(f"Invalid signature {signature} != {UNITY3D_SIGNATURE}") data = decrypt_key(self.key, self.data, DECRYPT_KEY) - data = to_uint4_array(data) + data = bytes( + nibble for byte in data for nibble in (byte >> 4, byte & 0xF) + ) self.index = data[:0x10] self.substitute = bytes( data[0x10 + i * 4 + j] for j in range(4) for i in range(4) ) def decrypt_block(self, data: bytes, index: int): - if UnityPyBoost: return UnityPyBoost.decrypt_block(self.index, self.substitute, data, index) @@ -119,7 +112,7 @@ def decrypt_block(self, data: bytes, index: int): index += 1 return data - def decrypt_byte(self, view: bytearray, offset: int, index: int): + def decrypt_byte(self, view: Union[bytearray, memoryview], offset: int, index: int): b = ( self.substitute[((index >> 2) & 3) + 4] + self.substitute[index & 3] @@ -133,7 +126,7 @@ def decrypt_byte(self, view: bytearray, offset: int, index: int): b = view[offset] return b, offset + 1, index + 1 - def decrypt(self, data: bytearray, index: int, remaining: int): + def decrypt(self, data: Union[bytearray, memoryview], index: int, remaining: int): offset = 0 curByte, offset, index = self.decrypt_byte(data, offset, index) diff --git a/UnityPy/helpers/CompressionHelper.py b/UnityPy/helpers/CompressionHelper.py index 9572d40a..16894173 100644 --- a/UnityPy/helpers/CompressionHelper.py +++ b/UnityPy/helpers/CompressionHelper.py @@ -1,11 +1,10 @@ +import brotli import gzip import lzma +import lz4.block import struct from typing import Tuple -import brotli -import lz4.block - GZIP_MAGIC: bytes = b"\x1f\x8b" BROTLI_MAGIC: bytes = b"brotli" diff --git a/UnityPy/helpers/ImportHelper.py b/UnityPy/helpers/ImportHelper.py index 7e7e2a4b..df4fc5f0 100644 --- a/UnityPy/helpers/ImportHelper.py +++ b/UnityPy/helpers/ImportHelper.py @@ -1,6 +1,7 @@ from __future__ import annotations + import os -from typing import Union, List +from typing import Union, List, Optional, Tuple from .CompressionHelper import BROTLI_MAGIC, GZIP_MAGIC from ..enums import FileType from ..streams import EndianBinaryReader @@ -18,7 +19,7 @@ def list_all_files(directory: str) -> List[str]: val for sublist in [ [os.path.join(dir_path, filename) for filename in filenames] - for (dir_path, dirn_ames, filenames) in os.walk(directory) + for (dir_path, dirnames, filenames) in os.walk(directory) if ".git" not in dir_path ] for val in sublist @@ -34,14 +35,14 @@ def find_all_files(directory: str, search_str: str) -> List[str]: for filename in filenames if search_str in filename ] - for (dir_path, dirn_ames, filenames) in os.walk(directory) + for (dir_path, dirnames, filenames) in os.walk(directory) if ".git" not in dir_path ] for val in sublist ] -def check_file_type(input_) -> Union[FileType, EndianBinaryReader]: +def check_file_type(input_) -> Tuple[Optional[FileType], Optional[EndianBinaryReader]]: if isinstance(input_, str) and os.path.isfile(input_): reader = EndianBinaryReader(open(input_, "rb")) elif isinstance(input_, EndianBinaryReader): @@ -124,10 +125,10 @@ def check_file_type(input_) -> Union[FileType, EndianBinaryReader]: def parse_file( reader: EndianBinaryReader, - parent, + parent: files.File, name: str, - typ: FileType = None, - is_dependency=False, + typ: Optional[FileType] = None, + is_dependency: bool = False ) -> Union[files.File, EndianBinaryReader]: if typ is None: typ, _ = check_file_type(reader) diff --git a/UnityPy/helpers/MeshHelper.py b/UnityPy/helpers/MeshHelper.py index e32f2584..83c7028a 100644 --- a/UnityPy/helpers/MeshHelper.py +++ b/UnityPy/helpers/MeshHelper.py @@ -1,9 +1,9 @@ from __future__ import annotations + import math import struct -from typing import Optional, List, Tuple, Union, TypeVar +from typing import Optional, List, Sequence, Tuple, Union, TypeVar -from ..enums.MeshTopology import MeshTopology from ..classes.generated import ( ChannelInfo, StreamInfo, @@ -13,8 +13,7 @@ Vector3f, Vector4f, ) -from .PackedBitVector import unpack_floats, unpack_ints - +from ..enums.MeshTopology import MeshTopology from ..enums.VertexFormat import ( VertexChannelFormat, VertexFormat2017, @@ -23,6 +22,7 @@ VERTEX_FORMAT_2017_STRUCT_TYPE_MAP, VERTEX_FORMAT_STRUCT_TYPE_MAP, ) +from .PackedBitVector import unpack_floats, unpack_ints from .ResourceReader import get_resource_data try: @@ -37,11 +37,11 @@ T = TypeVar("T") -def flat_list_to_tuples(data: List[T], item_size: int) -> List[tuple[T]]: +def flat_list_to_tuples(data: Sequence[T], item_size: int) -> List[tuple[T, ...]]: return [tuple(data[i : i + item_size]) for i in range(0, len(data), item_size)] -def vector_list_to_tuples(data: List[Vector2f, Vector3f, Vector4f]) -> List[tuple]: +def vector_list_to_tuples(data: List[Union[Vector2f, Vector3f, Vector4f]]) -> List[tuple]: if isinstance(data[0], Vector2f): return [(v.x, v.y) for v in data] elif isinstance(data[0], Vector3f): @@ -52,7 +52,7 @@ def vector_list_to_tuples(data: List[Vector2f, Vector3f, Vector4f]) -> List[tupl raise ValueError("Unknown vector type") -def zeros(shape: Tuple[int, ...]) -> list: +def zeros(shape: Union[Tuple[int], Tuple[int, int]]) -> Union[List, List[List]]: if len(shape) == 1: return [0] * shape[0] elif len(shape) == 2: @@ -97,7 +97,7 @@ def __init__( src: Union[Mesh, SpriteRenderData], version: Optional[Tuple[int, int, int, int]] = None, endianess: str = "<", - ) -> None: + ): self.src = src self.endianess = endianess if version is not None: diff --git a/UnityPy/helpers/PackedBitVector.py b/UnityPy/helpers/PackedBitVector.py index 8ad9ac16..13267091 100644 --- a/UnityPy/helpers/PackedBitVector.py +++ b/UnityPy/helpers/PackedBitVector.py @@ -1,6 +1,5 @@ from typing import Optional, List, TYPE_CHECKING, Tuple -# from ..objects.math import Quaternionf if TYPE_CHECKING: from ..classes.generated import PackedBitVector diff --git a/UnityPy/helpers/ResourceReader.py b/UnityPy/helpers/ResourceReader.py index e058b7ef..bf75ae03 100644 --- a/UnityPy/helpers/ResourceReader.py +++ b/UnityPy/helpers/ResourceReader.py @@ -1,51 +1,37 @@ import ntpath from ..streams import EndianBinaryReader -from ..files import File +from typing import TYPE_CHECKING +if TYPE_CHECKING: + from ..files.SerializedFile import SerializedFile -def get_resource_data(*args): - """ - Input: - Option 1: - 0 - path - file path - 1 - assets_file - SerializedFile - 2 - offset - - 3 - size - - Option 2: - 0 - reader - EndianBinaryReader - 1 - offset - - 2 - size - - -> -2 = offset, -1 = size - """ - if len(args) == 4: - res_path, assets_file, offset, size = args - basename = ntpath.basename(res_path) - name, ext = ntpath.splitext(basename) - possible_names = [ - basename, - f"{name}.resource", - f"{name}.assets.resS", - f"{name}.resS", - ] - environment = assets_file.environment - reader = None +def get_resource_data(res_path: str, assets_file: "SerializedFile", offset: int, size: int): + basename = ntpath.basename(res_path) + name, ext = ntpath.splitext(basename) + possible_names = [ + basename, + f"{name}.resource", + f"{name}.assets.resS", + f"{name}.resS", + ] + environment = assets_file.environment + reader = None + for possible_name in possible_names: + reader = environment.get_cab(possible_name) + if reader: + break + if not reader: + assets_file.load_dependencies(possible_names) for possible_name in possible_names: reader = environment.get_cab(possible_name) if reader: break if not reader: - assets_file.load_dependencies(possible_names) - for possible_name in possible_names: - reader = environment.get_cab(possible_name) - if reader: - break - if not reader: - raise FileNotFoundError(f"Resource file {basename} not found") - elif len(args) == 3: - reader, offset, size = args - else: - raise TypeError(f"3 or 4 arguments required, but only {len(args)} given") + raise FileNotFoundError(f"Resource file {basename} not found") + return _get_resource_data(reader, offset, size) + +def _get_resource_data(reader: EndianBinaryReader, offset: int, size: int): reader.Position = offset return reader.read_bytes(size) diff --git a/UnityPy/helpers/Tpk.py b/UnityPy/helpers/Tpk.py index a636d7d1..3b7dad02 100644 --- a/UnityPy/helpers/Tpk.py +++ b/UnityPy/helpers/Tpk.py @@ -4,7 +4,7 @@ from importlib.resources import open_binary from io import BytesIO from struct import Struct -from typing import Any, Dict, List, Tuple +from typing import Any, Dict, List, Optional, Tuple from .TypeTreeHelper import TypeTreeNode @@ -19,7 +19,9 @@ def init(): global TPKTYPETREE with BytesIO(data) as stream: - TPKTYPETREE = TpkFile(stream).GetDataBlob() + blob = TpkFile(stream).GetDataBlob() + assert isinstance(blob, TpkTypeTreeBlob) + TPKTYPETREE = blob def get_typetree_node(class_id: int, version: tuple): @@ -287,7 +289,7 @@ def __init__(self, stream: BytesIO) -> None: class UnityVersion(int): # https://github.com/AssetRipper/VersionUtilities/blob/master/VersionUtilities/UnityVersion.cs """ - use following static methos instead of the constructor(__init__): + use following static methods instead of the constructor(__init__): UnityVersion.fromStream(stream: BytesIO) UnityVersion.fromString(version: str) UnityVersion.fromList(major: int, minor: int, patch: int, build: int) @@ -300,7 +302,7 @@ def fromStream(stream: BytesIO) -> UnityVersion: @staticmethod def fromString(version: str) -> UnityVersion: - return UnityVersion(version.split(".")) + return UnityVersion.fromList(*map(int, version.split("."))) @staticmethod def fromList( @@ -338,8 +340,8 @@ class TpkUnityClass: Name: int Base: int Flags: TpkUnityClassFlags - EditorRootNode: int - ReleaseRootNode: int + EditorRootNode: Optional[int] + ReleaseRootNode: Optional[int] def __init__(self, stream: BytesIO) -> None: self.Name, self.Base, Flags = TpkUnityClass.Struct.unpack( @@ -473,7 +475,7 @@ def Count(self) -> int: class TpkCommonString: __slots__ = ("VersionInformation", "StringBufferIndices") VersionInformation: List[Tuple[UnityVersion, int]] - StringBufferIndices: List[int] + StringBufferIndices: Tuple[int] def __init__(self, stream: BytesIO) -> None: (versionCount,) = INT32.unpack(stream.read(INT32.size)) diff --git a/UnityPy/helpers/TypeTreeHelper.py b/UnityPy/helpers/TypeTreeHelper.py index 8f6e3659..22e3a66c 100644 --- a/UnityPy/helpers/TypeTreeHelper.py +++ b/UnityPy/helpers/TypeTreeHelper.py @@ -1,4 +1,4 @@ -from __future__ import annotations +from __future__ import annotations import re from typing import TYPE_CHECKING, Any, Optional, Union @@ -49,6 +49,7 @@ "string": EndianBinaryReader.read_aligned_string, "TypelessData": EndianBinaryReader.read_byte_array, } + FUNCTION_READ_MAP_ARRAY = { "SInt8": EndianBinaryReader.read_byte_array, "UInt8": EndianBinaryReader.read_u_byte_array, @@ -83,7 +84,7 @@ def copy(self) -> TypeTreeConfig: return TypeTreeConfig(self.as_dict, self.assetsfile, self.has_registry) -def get_ref_type_node(ref_object: dict, assetfile: SerializedFile) -> TypeTreeNode: +def get_ref_type_node(ref_object: dict, assetfile: SerializedFile) -> Optional[TypeTreeNode]: typ = ref_object["type"] if isinstance(typ, dict): cls = typ["class"] diff --git a/UnityPy/helpers/TypeTreeNode.py b/UnityPy/helpers/TypeTreeNode.py index 8f847b0d..22ea9bdd 100644 --- a/UnityPy/helpers/TypeTreeNode.py +++ b/UnityPy/helpers/TypeTreeNode.py @@ -2,7 +2,7 @@ import re from struct import Struct -from typing import TYPE_CHECKING, Dict, Iterator, List, Optional, Tuple, Any, Union +from typing import TYPE_CHECKING, Dict, Iterator, List, Optional, Tuple, Union from attrs import define, field