diff --git a/UnityPy/enums/BundleFile.py b/UnityPy/enums/BundleFile.py index fec1cf4a..be7dd85d 100644 --- a/UnityPy/enums/BundleFile.py +++ b/UnityPy/enums/BundleFile.py @@ -6,7 +6,8 @@ class CompressionFlags(IntFlag): LZMA = 1 LZ4 = 2 LZ4HC = 3 - LZHAM = 4 + COMPRESSION_4 = 4 + COMPRESSION_5 = 5 class ArchiveFlagsOld(IntFlag): diff --git a/UnityPy/files/BundleFile.py b/UnityPy/files/BundleFile.py index 3cd650b5..f4eb8681 100644 --- a/UnityPy/files/BundleFile.py +++ b/UnityPy/files/BundleFile.py @@ -1,6 +1,7 @@ # TODO: implement encryption for saving files from collections import namedtuple import re +import os from typing import Optional, Tuple, Union from . import File @@ -486,6 +487,17 @@ def save_web_raw(self, writer: EndianBinaryWriter): # Write compressed content writer.write(compressed_content) + @staticmethod + def read_long_length_no_check(ip: bytearray, pos: int) -> tuple[int, int]: + b = 0 + l = 0 + while True: + b = ip[pos] + pos += 1 + l += b + if b != 255: + break + return l, pos def decompress_data( self, @@ -516,8 +528,36 @@ def decompress_data( if self.decryptor is not None and flags & 0x100: compressed_data = self.decryptor.decrypt_block(compressed_data, index) return CompressionHelper.decompress_lz4(compressed_data, uncompressed_size) - elif comp_flag == CompressionFlags.LZHAM: # LZHAM - raise NotImplementedError("LZHAM decompression not implemented") + elif comp_flag in [CompressionFlags.COMPRESSION_4, CompressionFlags.COMPRESSION_5] and os.environ.get('UNITYPY_AK'): # ARKNIGHTS CUSTOM + ip = 0 + op = 0 + AK_LITERAL_LENGTH_MASK = ((1 << 4) - 1) & 0xFF + AK_MATCH_LENGTH_MASK = (~AK_LITERAL_LENGTH_MASK) & 0xFF + fixed_compressed_data = bytearray(compressed_data) + while True: + literal_length, match_length = ( + fixed_compressed_data[ip] & AK_LITERAL_LENGTH_MASK, + (fixed_compressed_data[ip] & AK_MATCH_LENGTH_MASK) >> 4 & 0xff, + ) + fixed_compressed_data[ip] = (literal_length << 4 | match_length) & 0xFF + ip += 1 + if literal_length == 15: + l, ip = BundleFile.read_long_length_no_check(fixed_compressed_data, ip) + literal_length += l + op += literal_length + ip += literal_length + if uncompressed_size == op: # MFLIMIT end of block + break + offset = fixed_compressed_data[ip + 1] | fixed_compressed_data[ip] << 8 + fixed_compressed_data[ip] = offset & 0xFF + fixed_compressed_data[ip + 1] = (offset >> 8) & 0xFF + ip += 2 + if match_length == 15: + m, ip = BundleFile.read_long_length_no_check(fixed_compressed_data, ip) + match_length += m + match_length += 4 # MINMATCH + op += match_length + return CompressionHelper.decompress_lz4(bytes(fixed_compressed_data), uncompressed_size) else: return compressed_data diff --git a/tests/samples/a1dd7f06fc870d5df6ccb483c7cd5686.bin b/tests/samples/a1dd7f06fc870d5df6ccb483c7cd5686.bin new file mode 100644 index 00000000..4b4c4d41 Binary files /dev/null and b/tests/samples/a1dd7f06fc870d5df6ccb483c7cd5686.bin differ diff --git a/tests/test_ak.py b/tests/test_ak.py new file mode 100644 index 00000000..38201472 --- /dev/null +++ b/tests/test_ak.py @@ -0,0 +1,29 @@ +import os +from tempfile import TemporaryDirectory + +from UnityPy.tools.extractor import extract_assets + +SAMPLES = os.path.join(os.path.dirname(os.path.abspath(__file__)), "samples") + + +def test_ak(): + os.environ['UNITYPY_AK'] = '1' + + temp_dir = TemporaryDirectory(prefix="unitypy_test_ak") + extract_assets( + SAMPLES, + temp_dir.name, + True, + ) + files = [ + os.path.relpath(os.path.join(root, f), temp_dir.name) + for root, dirs, files in os.walk(temp_dir.name) + for f in files + ] + print(files) + temp_dir.cleanup() + assert len(files) == 46 + + +if __name__ == "__main__": + test_ak()