Fix UnityCN Decryption + brute-force function

UnityPy/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "1.9.24"
+__version__ = "1.9.25"
 
 from .environment import Environment
 from .helpers.ArchiveStorageManager import set_assetbundle_decrypt_key

UnityPy/files/BundleFile.py

@@ -1,7 +1,7 @@
 # TODO: implement encryption for saving files
 from collections import namedtuple
 import re
-from typing import Tuple
+from typing import Tuple, Union
 
 from . import File
 from ..enums import ArchiveFlags, ArchiveFlagsOld, CompressionFlags
@@ -106,6 +106,9 @@ def read_fs(self, reader: EndianBinaryReader):
         else:
             self.dataflags = ArchiveFlags(self.dataflags)
 
+        if self.dataflags & self.dataflags.UsesAssetBundleEncryption:
+            self.decryptor = ArchiveStorageManager.ArchiveStorageDecryptor(reader)
+
         if self.version >= 7:
             reader.align_stream(16)
 
@@ -117,8 +120,6 @@ def read_fs(self, reader: EndianBinaryReader):
             blocksInfoBytes = reader.read_bytes(compressedSize)
             reader.Position = start
         else:  # 0x40 kArchiveBlocksAndDirectoryInfoCombined
-            if self.dataflags & self.dataflags.UsesAssetBundleEncryption:
-                self.decryptor = ArchiveStorageManager.ArchiveStorageDecryptor(reader)
             blocksInfoBytes = reader.read_bytes(compressedSize)
 
         blocksInfoBytes = self.decompress_data(
@@ -385,7 +386,11 @@ def save_fs(self, writer: EndianBinaryWriter, data_flag: int, block_info_flag: i
         writer.Position = writer_end_pos
 
     def decompress_data(
-        self, compressed_data: bytes, uncompressed_size: int, flags: int, index: int = 0
+        self,
+        compressed_data: bytes,
+        uncompressed_size: int,
+        flags: Union[int, ArchiveFlags, ArchiveFlagsOld],
+        index: int = 0,
     ) -> bytes:
         """
         Parameters
@@ -401,7 +406,7 @@ def decompress_data(
         -------
         bytes
             The decompressed data."""
-        comp_flag = flags & ArchiveFlags.CompressionTypeMask
+        comp_flag = CompressionFlags(flags & ArchiveFlags.CompressionTypeMask)
 
         if comp_flag == CompressionFlags.LZMA:  # LZMA
             return CompressionHelper.decompress_lzma(compressed_data)

UnityPy/helpers/ArchiveStorageManager.py

@@ -1,4 +1,5 @@
 # based on: https://github.com/Razmoth/PGRStudio/blob/master/AssetStudio/PGR/PGR.cs
+import re
 from typing import Tuple, Union
 
 from ..streams import EndianBinaryReader
@@ -20,7 +21,8 @@ def set_assetbundle_decrypt_key(key: Union[bytes, str]):
 
 def read_vector(reader: EndianBinaryReader) -> Tuple[bytes, bytes]:
     data = reader.read_bytes(0x10)
-    key = reader.read_string_to_null().encode("utf-8", "surrogateescape")
+    key = reader.read_bytes(0x10)
+    reader.Position += 1
 
     return data, key
 
@@ -32,6 +34,28 @@ def decrypt_key(key: bytes, data: bytes, keybytes: bytes):
     return bytes(x ^ y for x, y in zip(data, key))
 
 
+def brute_force_key(
+    fp: str,
+    key_sig: bytes,
+    data_sig: bytes,
+    pattern: re.Pattern = re.compile(rb"(?=(\w{16}))"),
+    verbose: bool = False,
+):
+    with open(fp, "rb") as f:
+        data = f.read()
+
+    matches = pattern.findall(data)
+    for i, key in enumerate(matches):
+        if verbose:
+            print(f"Trying {i + 1}/{len(matches)} - {key}")
+        signature = decrypt_key(key_sig, data_sig, key)
+        if signature == UNITY3D_SIGNATURE:
+            if verbose:
+                print(f"Found key: {key}")
+            return key
+    return None
+
+
 def to_uint4_array(source: bytes, offset: int = 0):
     buffer = bytearray(len(source) * 2)
     for j in range(len(source)):
@@ -46,16 +70,25 @@ class ArchiveStorageDecryptor:
     substitute: bytes = bytes(0x10)
 
     def __init__(self, reader: EndianBinaryReader) -> None:
-        if DECRYPT_KEY is None:
-            raise LookupError(
-                "The BundleFile is encrypted, but no key was provided!\nYou can set the key via UnityPy.set_assetbundle_decrypt_key(key)"
-            )
         self.unknown_1 = reader.read_u_int()
 
         # read vector data/key vectors
         self.data, self.key = read_vector(reader)
         self.data_sig, self.key_sig = read_vector(reader)
 
+        if DECRYPT_KEY is None:
+            raise LookupError(
+                "\n".join(
+                    [
+                        "The BundleFile is encrypted, but no key was provided!",
+                        "You can set the key via UnityPy.set_assetbundle_decrypt_key(key).",
+                        "To try brute-forcing the key, use UnityPy.helpers.ArchiveStorageManager.brute_force_key(fp, key_sig, data_sig)",
+                        f"with  key_sig = {self.key_sig}, data_sig = {self.data_sig},"
+                        "and fp being the path to global-metadata.dat or a memory dump.",
+                    ]
+                )
+            )
+
         signature = decrypt_key(self.key_sig, self.data_sig, DECRYPT_KEY)
         if signature != UNITY3D_SIGNATURE:
             raise Exception(f"Invalid signature {signature} != {UNITY3D_SIGNATURE}")