Improve config and files module

UnityPy/config.py

@@ -1,20 +1,41 @@
-# used when no version is defined by the SerializedFile or its BundleFile
-FALLBACK_UNITY_VERSION = "2.5.0f5"
-# determines if the typetree structures for the Object types will be parsed
-# disabling this will reduce the load time by a lot (half of the time is spend on parsing the typetrees)
-#  but it will also prevent saving an edited file
+import warnings
+
+from .exceptions import UnityVersionFallbackError, UnityVersionFallbackWarning
+
+
+FALLBACK_UNITY_VERSION = None
+"""The Unity version to use when no version is defined
+   by the SerializedFile or its BundleFile.
+
+   You may manually configure this value to a version string, e.g. `2.5.0f5`.
+"""
+
 SERIALIZED_FILE_PARSE_TYPETREE = True
+"""Determines if the typetree structures for the Object types will be parsed.
 
-# GLOBAL WARNING SUPPRESSION
-FALLBACK_VERSION_WARNED = False  # for FALLBACK_UNITY_VERSION
+   Disabling this will reduce the load time by a lot (half of the time is spend on parsing the typetrees),
+   but it will also prevent saving an edited file.
+"""
+
+
+# WARNINGS CONTROL
+warnings.simplefilter("once", UnityVersionFallbackWarning)
 
 
 # GET FUNCTIONS
 def get_fallback_version():
-    global FALLBACK_VERSION_WARNED
-    if not FALLBACK_VERSION_WARNED:
-        print(
-            f"Warning: 0.0.0 version found, defaulting to UnityPy.config.FALLBACK_UNITY_VERSION ({FALLBACK_UNITY_VERSION})"  # noqa: E501
+    global FALLBACK_UNITY_VERSION
+
+    if not isinstance(FALLBACK_UNITY_VERSION, str):
+        raise UnityVersionFallbackError(
+            "No valid Unity version found, and the fallback version is not correctly configured. "
+            + "Please explicitly set the value of UnityPy.config.FALLBACK_UNITY_VERSION."
         )
-        FALLBACK_VERSION_WARNED = True
+
+    warnings.warn(
+        f"No valid Unity version found, defaulting to UnityPy.config.FALLBACK_UNITY_VERSION ({FALLBACK_UNITY_VERSION})",  # noqa: E501
+        category=UnityVersionFallbackWarning,
+        stacklevel=2
+    )
+
     return FALLBACK_UNITY_VERSION

UnityPy/exceptions.py

@@ -1,5 +1,14 @@
 class TypeTreeError(Exception):
-    def __init__(self, message, nodes):            
-        # Call the base class constructor with the parameters it needs
+    def __init__(self, message, nodes):
+        super().__init__(message)
+        self.nodes = nodes
+
+
+class UnityVersionFallbackError(Exception):
+    def __init__(self, message):
+        super().__init__(message)
+
+
+class UnityVersionFallbackWarning(UserWarning):
+    def __init__(self, message):
         super().__init__(message)
-        self.nodes = nodes

UnityPy/files/BundleFile.py

@@ -1,10 +1,11 @@
 # TODO: implement encryption for saving files
 from collections import namedtuple
 import re
-from typing import Tuple, Union
+from typing import Optional, Tuple, Union
 
 from . import File
 from ..enums import ArchiveFlags, ArchiveFlagsOld, CompressionFlags
+from ..exceptions import UnityVersionFallbackError
 from ..helpers import ArchiveStorageManager, CompressionHelper
 from ..streams import EndianBinaryReader, EndianBinaryWriter
 
@@ -21,12 +22,12 @@ class BundleFile(File.File):
     signature: str
     version_engine: str
     version_player: str
-    dataflags: Tuple[ArchiveFlags, ArchiveFlagsOld]
-    decryptor: ArchiveStorageManager.ArchiveStorageDecryptor = None
+    dataflags: Union[ArchiveFlags, ArchiveFlagsOld]
+    decryptor: Optional[ArchiveStorageManager.ArchiveStorageDecryptor] = None
     _uses_block_alignment: bool = False
 
     def __init__(
-        self, reader: EndianBinaryReader, parent: File, name: str = None, **kwargs
+        self, reader: EndianBinaryReader, parent: File, name: Optional[str] = None, **kwargs
     ):
         super().__init__(parent=parent, name=name, **kwargs)
         signature = self.signature = reader.read_string_to_null()
@@ -92,7 +93,7 @@ def read_fs(self, reader: EndianBinaryReader):
         # header
         compressedSize = reader.read_u_int()
         uncompressedSize = reader.read_u_int()
-        self.dataflags = reader.read_u_int()
+        dataflagsValue = reader.read_u_int()
 
         version = self.get_version_tuple()
         # https://issuetracker.unity3d.com/issues/files-within-assetbundles-do-not-start-on-aligned-boundaries-breaking-patching-on-nintendo-switch
@@ -105,9 +106,9 @@ def read_fs(self, reader: EndianBinaryReader):
             or (version[0] == 2021 and version < (2021, 3, 2))
             or (version[0] == 2022 and version < (2022, 1, 1))
         ):
-            self.dataflags = ArchiveFlagsOld(self.dataflags)
+            self.dataflags = ArchiveFlagsOld(dataflagsValue)
         else:
-            self.dataflags = ArchiveFlags(self.dataflags)
+            self.dataflags = ArchiveFlags(dataflagsValue)
 
         if self.dataflags & self.dataflags.UsesAssetBundleEncryption:
             self.decryptor = ArchiveStorageManager.ArchiveStorageDecryptor(reader)
@@ -201,8 +202,8 @@ def save(self, packer=None):
                 original - uses the original flags
         """
         # file_header
-        #     signature    (string_to_null)
-        #     format        (int)
+        #     signature         (string_to_null)
+        #     format            (int)
         #     version_player    (string_to_null)
         #     version_engine    (string_to_null)
         writer = EndianBinaryWriter()
@@ -252,11 +253,11 @@ def save_fs(self, writer: EndianBinaryWriter, data_flag: int, block_info_flag: i
         # data_flag
 
         # header:
-        #     bundle_size        (long)
-        #     compressed_size    (int)
-        #     uncompressed_size    (int)
-        #     flag                (int)
-        #     ?padding?            (bool)
+        #     bundle_size       (long)
+        #     compressed_size   (int)
+        #     uncompressed_size (int)
+        #     flag              (int)
+        #     ?padding?         (bool)
         #   This will be written at the end,
         #   because the size can only be calculated after the data compression,
 
@@ -266,21 +267,21 @@ def save_fs(self, writer: EndianBinaryWriter, data_flag: int, block_info_flag: i
         #     *read compressed_size -> uncompressed_size
         #     0x10 offset
         #     *read blocks infos of the data stream
-        #     count            (int)
+        #     count                 (int)
         #     (
-        #         uncompressed_size(uint)
-        #         compressed_size (uint)
-        #         flag(short)
+        #         uncompressed_size (uint)
+        #         compressed_size   (uint)
+        #         flag              (short)
         #     )
         #     *decompression via info.flag & 0x3F
 
         #     *afterwards the file positions
-        #     file_count        (int)
+        #     file_count    (int)
         #     (
         #         offset    (long)
-        #         size        (long)
-        #         flag        (int)
-        #         name        (string_to_null)
+        #         size      (long)
+        #         flag      (int)
+        #         name      (string_to_null)
         #     )
 
         # file list & file data
@@ -523,6 +524,14 @@ def decompress_data(
     def get_version_tuple(self) -> Tuple[int, int, int]:
         """Returns the version as a tuple."""
         version = self.version_engine
-        if not version or version == "0.0.0":
-            version = config.get_fallback_version()
-        return tuple(map(int, reVersion.match(version).groups()))
+        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_))

UnityPy/files/File.py

@@ -13,7 +13,7 @@
 DirectoryInfo = namedtuple("DirectoryInfo", "path offset size")
 
 
-class File(object):
+class File:
     name: str
     files: Dict[str, File]
     environment: Environment
@@ -34,7 +34,7 @@ def __init__(
         self.is_changed = False
         self.cab_file = "CAB-UnityPy_Mod.resS"
         self.parent = parent
-        self.environment = self.environment = (
+        self.environment = (
             getattr(parent, "environment", parent) if parent else None
         )
         self.name = basename(name) if isinstance(name, str) else ""
@@ -95,10 +95,10 @@ def read_files(self, reader: EndianBinaryReader, files: list):
             f.flags = getattr(node, "flags", 0)
             self.files[name] = f
 
-    def get_writeable_cab(self, name: str = None):
+    def get_writeable_cab(self, name: Optional[str] = None):
         """
         Creates a new cab file in the bundle that contains the given data.
-        This is usefull for asset types that use resource files.
+        This is useful for asset types that use resource files.
         """
 
         if not name:

UnityPy/files/ObjectReader.py

@@ -152,9 +152,11 @@ def write(
             writer.write_u_short(self.class_id)
 
         if header.version < 11:
+            assert self.is_destroyed is not None
             writer.write_u_short(self.is_destroyed)
 
         if 11 <= header.version < 17:
+            assert self.serialized_type is not None
             writer.write_short(self.serialized_type.script_type_index)
 
         if header.version == 15 or header.version == 16:
@@ -254,7 +256,7 @@ def save_typetree(
         self,
         tree: dict,
         nodes: Optional[Union[TypeTreeNode, List[dict[str, Union[str, int]]]]] = None,
-        writer: EndianBinaryWriter = None,
+        writer: Optional[EndianBinaryWriter] = None,
     ):
         node = self._get_typetree_node(nodes)
         if not writer:

UnityPy/files/SerializedFile.py

@@ -2,7 +2,7 @@
 
 import re
 from ntpath import basename
-from typing import TYPE_CHECKING, Dict, Generator, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Dict, Generator, Iterator, List, Optional, Tuple
 
 from attrs import define
 
@@ -24,7 +24,7 @@ class SerializedFileHeader:
     file_size: int
     version: int
     data_offset: int
-    endian: bytes
+    endian: str
     reserved: bytes
 
     def __init__(self, reader: EndianBinaryReader):
@@ -83,8 +83,13 @@ def __init__(self, header: SerializedFileHeader, reader: EndianBinaryReader):
 
     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)
@@ -110,7 +115,7 @@ def IsPatch(self):
 class SerializedType:
     class_id: int
     is_stripped_type: Optional[bool] = None
-    script_type_index: Optional[int] = -1
+    script_type_index: int = -1
     script_id: Optional[bytes] = None  # Hash128
     old_type_hash: Optional[bytes] = None  # Hash128
     node: Optional[TypeTreeNode] = None
@@ -119,7 +124,7 @@ class SerializedType:
     m_NameSpace: Optional[str] = None
     m_AssemblyName: Optional[str] = None
     # 21+
-    type_dependencies: Optional[List[int]] = None
+    type_dependencies: Optional[Tuple[int, ...]] = None
 
     def __init__(
         self,
@@ -170,9 +175,11 @@ def write(
         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:
@@ -181,7 +188,9 @@ def write(
                 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:
@@ -193,14 +202,20 @@ def write(
 
             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) -> Union[TypeTreeNode, None]:
+    def nodes(self) -> Optional[TypeTreeNode]:
         # for compatibility with old versions
         return self.node
 
@@ -221,8 +236,7 @@ class SerializedFile(File.File):
     _m_target_platform: int
     big_id_enabled: int
     userInformation: Optional[str]
-    assetbundle: AssetBundle
-    container: ContainerHelper
+    assetbundle: Optional[AssetBundle]
     _cache: Dict[str, Object]
 
     @property
@@ -399,7 +413,7 @@ class FileIdentifierFake:
 
         return cab
 
-    def save(self, packer: str = None) -> bytes:
+    def save(self, packer: Optional[str] = None) -> bytes:
         # 1. header -> has to be delayed until the very end
         # 2. data -> types, objects, scripts, ...
 
@@ -443,11 +457,13 @@ def save(self, packer: str = None) -> bytes:
             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
@@ -550,7 +566,7 @@ def __setitem__(self, key, value) -> None:
     def __delitem__(self, key) -> None:
         raise NotImplementedError("Deleting from the container is not allowed!")
 
-    def __iter__(self) -> Generator[str, None, None]:
+    def __iter__(self) -> Iterator[str]:
         return iter(self.keys())
 
     def __len__(self) -> int:
@@ -559,9 +575,6 @@ def __len__(self) -> int:
     def __getattr__(self, name: str) -> PPtr[Object]:
         return self.container_dict[name]
 
-    def __or__(self, other: ContainerHelper):
-        return ContainerHelper(list(set(self.container + other.container)))
-
     def __str__(self) -> str:
         return f'{{{", ".join(f"{key}: {value}" for key, value in self.items())}}}'