Improve helpers module

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)

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"
 

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)

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:

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
 

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)

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))

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"]

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