feat(UnityVersion)!: use a reimplemented UnityVersion class for Unity files and internal version handling, deprecated file.build_type

Author: K0lb3

UnityPy/files/BundleFile.py

@@ -1,12 +1,12 @@
 # TODO: implement encryption for saving files
 import re
 from collections import namedtuple
-from typing import Optional, Tuple, Union, cast
+from typing import Optional, Union, cast
 
 from .. import config
 from ..enums import ArchiveFlags, ArchiveFlagsOld, CompressionFlags
-from ..exceptions import UnityVersionFallbackError
 from ..helpers import ArchiveStorageManager, CompressionHelper
+from ..helpers.UnityVersion import UnityVersion
 from ..streams import EndianBinaryReader, EndianBinaryWriter
 from . import File
 
@@ -98,7 +98,7 @@ def read_fs(self, reader: EndianBinaryReader):
         uncompressedSize = reader.read_u_int()
         dataflagsValue = reader.read_u_int()
 
-        version = self.get_version_tuple()
+        version = self.parse_version()
         # https://issuetracker.unity3d.com/issues/files-within-assetbundles-do-not-start-on-aligned-boundaries-breaking-patching-on-nintendo-switch
         # Unity CN introduced encryption before the alignment fix was introduced.
         # Unity CN used the same flag for the encryption as later on the alignment fix,
@@ -513,17 +513,17 @@ def decompress_data(
         else:
             raise ValueError(f"Unknown compression! flag: {flags}, compression flag: {comp_flag.value}")
 
-    def get_version_tuple(self) -> Tuple[int, int, int]:
+    def parse_version(self) -> UnityVersion:
         """Returns the version as a tuple."""
-        version = self.version_engine
-        match = None
-        if version and version != "0.0.0":
-            match = reVersion.match(version)
-            if not match or len(match.groups()) < 3:
-                match = None
-        if not match:
-            match = reVersion.match(config.get_fallback_version())
-            if not match or len(match.groups()) < 3:
-                raise UnityVersionFallbackError("Illegal fallback version format")
-        map_ = map(int, match.groups())
-        return (next(map_), next(map_), next(map_))
+        version = None
+        version_str = self.version_engine
+        try:
+            version = UnityVersion.from_str(version_str)
+        except ValueError:
+            pass
+
+        if version is None or version.major == 0:
+            version_str = config.get_fallback_version()
+            version = UnityVersion.from_str(version_str)
+
+        return version

UnityPy/files/ObjectReader.py

@@ -21,10 +21,11 @@
 from ..helpers import TypeTreeHelper
 from ..helpers.Tpk import get_typetree_node
 from ..helpers.TypeTreeNode import TypeTreeNode
+from ..helpers.UnityVersion import UnityVersion
 from ..streams import EndianBinaryReader, EndianBinaryWriter
 
 if TYPE_CHECKING:
-    from ..files.SerializedFile import BuildType, SerializedFile, SerializedType
+    from ..files.SerializedFile import SerializedFile, SerializedType
 
 T = TypeVar("T")
 NodeInput = Union[TypeTreeNode, List[Dict[str, Union[str, int]]]]
@@ -34,10 +35,9 @@ class ObjectReader(Generic[T]):
     assets_file: SerializedFile
     reader: EndianBinaryReader
     data: bytes
-    version: Tuple[int, int, int, int]
+    version: UnityVersion
     version2: int
     platform: BuildTarget
-    build_type: BuildType
     path_id: int
     byte_start_offset: Tuple[int, int]
     byte_start: int
@@ -61,7 +61,6 @@ def __init__(self, assets_file: SerializedFile, reader: EndianBinaryReader):
         self.version = assets_file.version
         self.version2 = assets_file.header.version
         self.platform = assets_file.target_platform
-        self.build_type = assets_file.build_type
 
         header = assets_file.header
         types = assets_file.types

UnityPy/files/SerializedFile.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import re
 from ntpath import basename
 from typing import TYPE_CHECKING, Dict, List, Optional, Tuple
 
@@ -10,6 +9,7 @@
 from ..enums import BuildTarget, ClassIDType, CommonString
 from ..helpers.ContainerHelper import ContainerHelper
 from ..helpers.TypeTreeHelper import TypeTreeNode
+from ..helpers.UnityVersion import UnityVersion
 from ..streams import EndianBinaryWriter
 from . import BundleFile, File, ObjectReader
 
@@ -93,22 +93,6 @@ def write(self, header: SerializedFileHeader, writer: EndianBinaryWriter):
         writer.write_string_to_null(self.path)
 
 
-@define(slots=True)
-class BuildType:
-    build_type: str
-
-    def __init__(self, build_type):
-        self.build_type = build_type
-
-    @property
-    def IsAlpha(self):
-        return self.build_type == "a"
-
-    @property
-    def IsPatch(self):
-        return self.build_type == "p"
-
-
 @define(slots=True, init=False)
 class SerializedType:
     class_id: int
@@ -220,9 +204,8 @@ def nodes(self) -> Optional[TypeTreeNode]:
 
 class SerializedFile(File.File):
     reader: EndianBinaryReader
-    version: Tuple[int, int, int, int]
+    version: UnityVersion
     unity_version: str
-    build_type: BuildType
     target_platform: BuildTarget
     _enable_type_tree: bool
     types: List[SerializedType]
@@ -253,8 +236,6 @@ def __init__(self, reader: EndianBinaryReader, parent=None, name=None, **kwargs)
         self.reader = reader
 
         self.unity_version = "2.5.0f5"
-        self.version = (0, 0, 0, 0)
-        self.build_type = BuildType("")
         self.target_platform = BuildTarget.UnknownPlatform
         self._enable_type_tree = True
         self.types = []
@@ -375,10 +356,7 @@ def set_version(self, string_version: str):
                 string_version = self.parent.version_engine
             if not string_version or string_version == "0.0.0":
                 string_version = config.get_fallback_version()
-        build_type = re.findall(r"([^\d.])", string_version)
-        self.build_type = BuildType(build_type[0] if build_type else "")
-        version_split = re.split(r"\D", string_version)
-        self.version = tuple(int(x) for x in version_split[:4])
+        self.version = UnityVersion.from_str(string_version)
 
     def get_writeable_cab(self, name: str = "CAB-UnityPy_Mod.resS"):
         """

UnityPy/helpers/Tpk.py

@@ -4,12 +4,15 @@
 from importlib.resources import open_binary
 from io import BytesIO
 from struct import Struct
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple, TypeVar
 
 from .TypeTreeHelper import TypeTreeNode
+from .UnityVersion import UnityVersion
 
-TPKTYPETREE: TpkTypeTreeBlob = None
-CLASSES_CACHE: Dict[Tuple[int, tuple], TypeTreeNode] = {}
+T = TypeVar("T")
+
+TPKTYPETREE: TpkTypeTreeBlob = None  # pyright: ignore[reportAssignmentType]
+CLASSES_CACHE: Dict[Tuple[int, UnityVersion], TypeTreeNode] = {}
 NODES_CACHE: Dict[TpkUnityClass, TypeTreeNode] = {}
 
 
@@ -24,14 +27,14 @@ def init():
         TPKTYPETREE = blob
 
 
-def get_typetree_node(class_id: int, version: tuple):
+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(UnityVersion.fromList(*version))
+    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))
 
@@ -87,15 +90,6 @@ class TpkCompressionType(IntEnum):
     Brotli = 3
 
 
-class UnityVersionType(IntEnum):
-    Alpha = 0
-    Beta = 1
-    China = 2
-    Final = 3
-    Patch = 4
-    Experimental = 5
-
-
 class TpkDataType(IntEnum):
     TypeTreeInformation = 0
     Collection = 1
@@ -168,7 +162,7 @@ def __init__(self, stream: BytesIO):
             raise Exception("Invalid compressed size")
 
     def GetDataBlob(self) -> TpkDataBlob:
-        decompressed = None
+        decompressed: bytes
         if self.CompressionType == TpkCompressionType.NONE:
             decompressed = self.CompressedBytes
 
@@ -185,7 +179,7 @@ def GetDataBlob(self) -> TpkDataBlob:
         elif self.CompressionType == TpkCompressionType.Brotli:
             import brotli
 
-            decompressed: bytes = brotli.decompress(self.CompressedBytes)
+            decompressed = brotli.decompress(self.CompressedBytes)
 
         else:
             raise Exception("Invalid compression type")
@@ -228,7 +222,7 @@ class TpkTypeTreeBlob(TpkDataBlob):
     def __init__(self, stream: BytesIO) -> None:
         (self.CreationTime,) = INT64.unpack(stream.read(INT64.size))
         (versionCount,) = INT32.unpack(stream.read(INT32.size))
-        self.Versions = [UnityVersion.fromStream(stream) for _ in range(versionCount)]
+        self.Versions = [read_version(stream) for _ in range(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)
@@ -282,52 +276,6 @@ def __init__(self, stream: BytesIO) -> None:
 ######################################################################################
 
 
-class UnityVersion(int):
-    # https://github.com/AssetRipper/VersionUtilities/blob/master/VersionUtilities/UnityVersion.cs
-    """
-    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)
-    """
-
-    @staticmethod
-    def fromStream(stream: BytesIO) -> UnityVersion:
-        (m_data,) = UINT64.unpack(stream.read(UINT64.size))
-        return UnityVersion(m_data)
-
-    @staticmethod
-    def fromString(version: str) -> UnityVersion:
-        return UnityVersion.fromList(*map(int, version.split(".")))
-
-    @staticmethod
-    def fromList(major: int = 0, minor: int = 0, patch: int = 0, build: int = 0) -> UnityVersion:
-        return UnityVersion(major << 48 | minor << 32 | patch << 16 | build)
-
-    @property
-    def major(self) -> int:
-        return (self >> 48) & 0xFFFF
-
-    @property
-    def minor(self) -> int:
-        return (self >> 32) & 0xFFFF
-
-    @property
-    def build(self) -> int:
-        return (self >> 16) & 0xFFFF
-
-    @property
-    def type(self) -> int:
-        return UnityVersionType(self >> 8) & 0xFF
-
-    @property
-    def type_number(self) -> int:
-        return self & 0xFF
-
-    def __repr__(self) -> str:
-        return f"UnityVersion {self.major}.{self.minor}.{self.build}.{self.type_number}"
-
-
 class TpkUnityClass:
     __slots__ = ("Name", "Base", "Flags", "EditorRootNode", "ReleaseRootNode")
     Struct = Struct("<HHb")
@@ -374,20 +322,20 @@ class TpkClassInformation:
     __slots__ = ("ID", "Classes")
     ID: int
     # TODO - might want to use dict
-    Classes: List[Tuple[UnityVersion, TpkUnityClass]]
+    Classes: List[Tuple[UnityVersion, Optional[TpkUnityClass]]]
 
     def __init__(self, stream: BytesIO) -> None:
         (self.ID,) = INT32.unpack(stream.read(INT32.size))
         (count,) = INT32.unpack(stream.read(INT32.size))
         self.Classes = [
             (
-                UnityVersion.fromStream(stream),
+                read_version(stream),
                 TpkUnityClass(stream) if stream.read(1)[0] else None,
             )
             for _ in range(count)
         ]
 
-    def getVersionedClass(self, version: UnityVersion) -> TpkUnityClass:
+    def getVersionedClass(self, version: UnityVersion) -> Optional[TpkUnityClass]:
         return get_item_for_version(version, self.Classes)
 
 
@@ -469,7 +417,7 @@ class TpkCommonString:
 
     def __init__(self, stream: BytesIO) -> None:
         (versionCount,) = INT32.unpack(stream.read(INT32.size))
-        self.VersionInformation = [(UnityVersion.fromStream(stream), stream.read(1)[0]) for _ in range(versionCount)]
+        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))
@@ -512,7 +460,11 @@ def read_data(stream: BytesIO) -> bytes:
     return stream.read(INT32.unpack(stream.read(INT32.size))[0])
 
 
-def get_item_for_version(exactVersion: UnityVersion, items: List[Tuple[UnityVersion, Any]]) -> Any:
+def read_version(stream: BytesIO) -> UnityVersion:
+    return UnityVersion(UINT64.unpack(stream.read(UINT64.size))[0])
+
+
+def get_item_for_version(exactVersion: UnityVersion, items: List[Tuple[UnityVersion, T]]) -> T:
     ret = None
     for version, item in items:
         if exactVersion >= version: