from . import File

from ..helpers import CompressionHelper

from ..streams import EndianBinaryReader, EndianBinaryWriter

from collections import namedtuple

BlockInfo = namedtuple("BlockInfo", "uncompressedSize compressedSize flags")

DirectoryInfoFS = namedtuple("DirectoryInfoFS", "offset size flags path")

class BundleFile(File.File):

format: int

is_changed: bool

signature: str

version_engine: str

version_player: str

def __init__(self, reader: EndianBinaryReader, parent: File, name: str=None):

super().__init__(parent=parent, name=name)

signature = self.signature = reader.read_string_to_null()

self.version = reader.read_u_int()

self.version_player = reader.read_string_to_null()

self.version_engine = reader.read_string_to_null()

if signature == "UnityArchive":

raise NotImplemented("BundleFile - UnityArchive")

elif signature in ["UnityWeb", "UnityRaw"]:

m_DirectoryInfo, blocksReader = self.read_web_raw(reader)

elif signature == "UnityFS":

m_DirectoryInfo, blocksReader = self.read_fs(reader)

else:

raise NotImplemented(f"Unknown Bundle signature: {signature}")

self.read_files(blocksReader, m_DirectoryInfo)

def read_web_raw(self, reader: EndianBinaryReader):

# def read_header_and_blocks_info(self, reader:EndianBinaryReader):

version = self.version

if version >= 4:

_hash = reader.read_bytes(16)

crc = reader.read_u_int()

minimumStreamedBytes = reader.read_u_int()

headerSize = reader.read_u_int()

numberOfLevelsToDownloadBeforeStreaming = reader.read_u_int()

levelCount = reader.read_int()

reader.Position += 4 * 2 * (levelCount - 1)

compressedSize = reader.read_u_int()

uncompressedSize = reader.read_u_int()

if version >= 2:

completeFileSize = reader.read_u_int()

if version >= 3:

fileInfoHeaderSize = reader.read_u_int()

reader.Position = headerSize

uncompressedBytes = CompressionHelper.decompress_lzma(

reader.read_bytes(compressedSize)

)

blocksReader = EndianBinaryReader(uncompressedBytes, offset=headerSize)

nodesCount = blocksReader.read_int()

m_DirectoryInfo = [

File.DirectoryInfo(

blocksReader.read_string_to_null(), # path

blocksReader.read_u_int(), # offset

blocksReader.read_u_int(), # size

)

for _ in range(nodesCount)

]

return m_DirectoryInfo, blocksReader

def read_fs(self, reader: EndianBinaryReader):

version = self.version

size = reader.read_long()

# header

compressedSize = reader.read_u_int()

uncompressedSize = reader.read_u_int()

self._data_flags = reader.read_u_int()

if self.version >= 7:

reader.align_stream(16)

_position = reader.Position

if self._data_flags & 0x80 != 0: # kArchiveBlocksInfoAtTheEnd

reader.Position = reader.Length - compressedSize

blocksInfoBytes = reader.read_bytes(compressedSize)

reader.Position = _position

else: # 0x40 kArchiveBlocksAndDirectoryInfoCombined

blocksInfoBytes = reader.read_bytes(compressedSize)

switch = self._data_flags & 0x3F

if switch == 1: # LZMA

blocksInfoBytes = CompressionHelper.decompress_lzma(blocksInfoBytes)

elif switch in [2, 3]: # LZ4, LZ4HC

blocksInfoBytes = CompressionHelper.decompress_lz4(

blocksInfoBytes, uncompressedSize

)

# elif switch == 4: #LZHAM:

blocksInfoReader = EndianBinaryReader(blocksInfoBytes, offset=_position)

uncompressedDataHash = blocksInfoReader.read_bytes(16)

blocksInfoCount = blocksInfoReader.read_int()

m_BlocksInfo = [

BlockInfo(

blocksInfoReader.read_u_int(), # uncompressedSize

blocksInfoReader.read_u_int(), # compressedSize

blocksInfoReader.read_u_short(), # flags

)

for _ in range(blocksInfoCount)

]

nodesCount = blocksInfoReader.read_int()

m_DirectoryInfo = [

DirectoryInfoFS(

blocksInfoReader.read_long(), # offset

blocksInfoReader.read_long(), # size

blocksInfoReader.read_u_int(), # flags

blocksInfoReader.read_string_to_null(), # path

)

for _ in range(nodesCount)

]

# def read_blocks(reader : EndianBinaryReader, blocksStream):

def decompress_block(blockInfo):

switch = blockInfo.flags & 0x3F # kStorageBlockCompressionTypeMask

if switch == 1: # LZMA

return CompressionHelper.decompress_lzma(

reader.read_bytes(blockInfo.compressedSize)

)

elif switch in [2, 3]: # LZ4, LZ4HC

return CompressionHelper.decompress_lz4(

reader.read_bytes(blockInfo.compressedSize),

blockInfo.uncompressedSize,

)

# elif switch == 4: #LZHAM:

else: # no compression

return reader.read_bytes(blockInfo.uncompressedSize)

if m_BlocksInfo:

self._block_info_flags = m_BlocksInfo[0].flags

blocksReader = EndianBinaryReader(

b"".join(decompress_block(blockInfo) for blockInfo in m_BlocksInfo),

offset=(blocksInfoReader.real_offset())

)

return m_DirectoryInfo, blocksReader

def save(self, packer=None):

"""

Rewrites the BundleFile and returns it as bytes object.

packer:

can be either one of the following strings

or tuple consisting of (block_info_flag, data_flag)

allowed strings:

none - no compression, default, safest bet

lz4 - lz4 compression

original - uses the original flags

"""

# file_header

# signature (string_to_null)

# format (int)

# version_player (string_to_null)

# version_engine (string_to_null)

writer = EndianBinaryWriter()

writer.write_string_to_null(self.signature)

writer.write_u_int(self.version)

writer.write_string_to_null(self.version_player)

writer.write_string_to_null(self.version_engine)

if self.signature == "UnityArchive":

raise NotImplemented("BundleFile - UnityArchive")

elif self.signature in ["UnityWeb", "UnityRaw"]:

raise NotImplemented("Saving Unity Web and Raw bundles isn't supported yet")

# self.save_web_raw(writer)

elif self.signature == "UnityFS":

if not packer or packer == "none":

self.save_fs(writer, 64, 64)

elif packer == "original":

self.save_fs(writer, block_info_flag=self._block_info_flags, data_flag=self._data_flags)

elif packer == "lz4":

self.save_fs(writer, block_info_flag=194, data_flag=2)

elif isinstance(packer, tuple):

self.save_fs(writer, *packer)

else:

raise NotImplemented("UnityFS - Packer:", packer)

return writer.bytes

def save_fs(self, writer: EndianBinaryWriter, block_info_flag : int, data_flag : int):

# header

# compressed blockinfo (block details & directionary)

# compressed assets

# 0b1000000 / 0b11000000 | 64 / 192 - uncompressed

# 0b11000010 | 194 - lz4

# block_info_flag

# 0 / 0b1000000 | 0 / 64 - uncompressed

# 0b1 | 1 - lzma

# 0b10 | 2 - lz4

# 0b11 | 3 - lz4hc [not implemented]

# 0b100 | 4 - lzham [not implemented]

# data_flag

# header:

# 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,

# block_info:

# *flag & 0x80 ? at the end : right after header

# *decompression via flag & 0x3F

# *read compressed_size -> uncompressed_size

# 0x10 offset

# *read blocks infos of the data stream

# count (int)

# (

# uncompressed_size(uint)

# compressed_size (uint)

# flag(short)

# )

# *decompression via info.flag & 0x3F

# *afterwards the file positions

# file_count (int)

# (

# offset (long)

# size (long)

# flag (int)

# name (string_to_null)

# )

# file list & file data

# prep nodes and build up block data

data_writer = EndianBinaryWriter()

files = [

(

name,

f.flags,

data_writer.write_bytes(

f.bytes if isinstance(f, (EndianBinaryReader, EndianBinaryWriter)) else f.save()

),

)

for name, f in self.files.items()

]

file_data = data_writer.bytes

data_writer.dispose()

uncompressed_data_size = len(file_data)

# compress the data

switch = data_flag & 0x3F

if switch == 1: # LZMA

file_data = CompressionHelper.compress_lzma(file_data)

elif switch in [2, 3]: # LZ4, LZ4HC

file_data = CompressionHelper.compress_lz4(file_data)

elif switch == 4: # LZHAM

raise NotImplementedError

# else no compression - data stays the same

compressed_data_size = len(file_data)

# write the block_info

# uncompressedDataHash

block_writer = EndianBinaryWriter(b"\x00" * 0x10)

# data block info

# block count

block_writer.write_int(1)

# uncompressed size

block_writer.write_u_int(uncompressed_data_size)

# compressed size

block_writer.write_u_int(compressed_data_size)

# flag

block_writer.write_u_short(data_flag)

# file block info

# file count

block_writer.write_int(len(files))

offset = 0

for f_name, f_flag, f_len in files:

# offset

block_writer.write_long(offset)

# size

block_writer.write_long(f_len)

offset += f_len

# flag

block_writer.write_u_int(f_flag)

# name

block_writer.write_string_to_null(f_name)

# compress the block data

block_data = block_writer.bytes

block_writer.dispose()

uncompressed_block_data_size = len(block_data)

switch = block_info_flag & 0x3F

if switch == 1: # LZMA

block_data = CompressionHelper.compress_lzma(block_data)

elif switch in [2, 3]: # LZ4, LZ4HC

block_data = CompressionHelper.compress_lz4(block_data)

elif switch == 4: # LZHAM

raise NotImplementedError

compressed_block_data_size = len(block_data)

# write the header info

## file size

writer.write_long(

writer.Length

+ 8

+ 4

+ 4

+ 4

+ (8 if self.version >= 7 else 0)

+ compressed_block_data_size

+ compressed_data_size

)

# compressed blockInfoBytes size

writer.write_u_int(compressed_block_data_size)

# uncompressed size

writer.write_u_int(uncompressed_block_data_size)

# compression flag

writer.write_u_int(block_info_flag)

if self.version >= 7:

# UnityFS\x00 - 8

# size 8

# comp sizes 4+4

# flag 4

# sum : 28 -> +8 alignment

writer.align_stream(16)

if (block_info_flag & 0x80) != 0: # at end of file

writer.write(file_data)

writer.write(block_data)

else:

writer.write(block_data)

writer.write(file_data)