Includes earlier uncommitted work: Tomie and Unity adapters, benchmark, structured output, call metrics. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
192 lines
7.5 KiB
Python
192 lines
7.5 KiB
Python
"""Bounded-memory UnityFS I/O for this Unity 2019 game.
|
|
|
|
Uncompressed media stays in a temporary disk file. Asset metadata is still read
|
|
and written by UnityPy; only the outer bundle and resource streams are spooled.
|
|
"""
|
|
from contextlib import contextmanager
|
|
import io
|
|
import hashlib
|
|
from pathlib import Path
|
|
import struct
|
|
import tempfile
|
|
import zlib
|
|
|
|
|
|
def temporary_file():
|
|
directory = Path(__file__).resolve().parent.parent / "data/tomie/tmp"
|
|
directory.mkdir(parents=True, exist_ok=True)
|
|
return tempfile.TemporaryFile(dir=directory)
|
|
|
|
|
|
class FileSlice(io.RawIOBase):
|
|
def __init__(self, backing, offset, size):
|
|
self.backing, self.offset, self.size, self.position = backing, offset, size, 0
|
|
|
|
def tell(self):
|
|
return self.position
|
|
|
|
def seek(self, offset, whence=0):
|
|
position = offset if whence == 0 else self.position + offset if whence == 1 else self.size + offset
|
|
if not 0 <= position <= self.size:
|
|
raise ValueError("Resource seek outside its bounds")
|
|
self.position = position
|
|
return position
|
|
|
|
def read(self, size=-1):
|
|
size = self.size - self.position if size < 0 else min(size, self.size - self.position)
|
|
self.backing.seek(self.offset + self.position)
|
|
result = self.backing.read(size)
|
|
self.position += len(result)
|
|
return result
|
|
|
|
|
|
def read_fs(bundle, reader):
|
|
from UnityPy.enums.BundleFile import ArchiveFlagsOld
|
|
from UnityPy.files.BundleFile import DirectoryInfoFS
|
|
from UnityPy.streams import EndianBinaryReader
|
|
bundle.size = reader.read_long()
|
|
compressed, uncompressed, flags = reader.read_u_int(), reader.read_u_int(), reader.read_u_int()
|
|
if bundle.parse_version()[:2] != (2019, 4) or flags & 0x200:
|
|
raise ValueError("Streaming bundle reader only supports unencrypted Unity 2019.4 assets")
|
|
bundle.dataflags = ArchiveFlagsOld(flags)
|
|
bundle._uses_block_alignment = True
|
|
reader.align_stream(16)
|
|
start = reader.Position
|
|
if flags & 0x80:
|
|
reader.Position = reader.Length - compressed
|
|
info = EndianBinaryReader(bundle.decompress_data(reader.read_bytes(compressed), uncompressed, bundle.dataflags))
|
|
if flags & 0x80:
|
|
reader.Position = start
|
|
info.read_bytes(16)
|
|
blocks = [(info.read_u_int(), info.read_u_int(), info.read_u_short()) for _ in range(info.read_int())]
|
|
nodes = [DirectoryInfoFS(info.read_long(), info.read_long(), info.read_u_int(), info.read_string_to_null())
|
|
for _ in range(info.read_int())]
|
|
bundle._block_info_flags = blocks[0][2] if blocks else 0
|
|
backing = temporary_file()
|
|
crc = 0
|
|
try:
|
|
for index, (raw_size, packed_size, block_flags) in enumerate(blocks):
|
|
data = bundle.decompress_data(reader.read_bytes(packed_size), raw_size, block_flags, index)
|
|
if len(data) != raw_size:
|
|
raise ValueError("Decompressed block size mismatch")
|
|
crc = zlib.crc32(data, crc)
|
|
backing.write(data)
|
|
backing.seek(0)
|
|
bundle.stream_crc = crc & 0xffffffff
|
|
return nodes, EndianBinaryReader(backing)
|
|
except BaseException:
|
|
backing.close()
|
|
raise
|
|
|
|
|
|
@contextmanager
|
|
def streaming_unity():
|
|
from UnityPy.files.BundleFile import BundleFile
|
|
from UnityPy.files.File import File
|
|
from UnityPy.files.SerializedFile import SerializedFile
|
|
from UnityPy.helpers import ImportHelper
|
|
from UnityPy.streams import EndianBinaryReader
|
|
old_fs, old_files = BundleFile.read_fs, File.read_files
|
|
|
|
def read_files(parent, reader, nodes):
|
|
for node in nodes:
|
|
name = node.path
|
|
if name.endswith((".resS", ".resource")) and hasattr(reader, "stream"):
|
|
source = EndianBinaryReader(FileSlice(reader.stream, node.offset, node.size))
|
|
else:
|
|
reader.Position = node.offset
|
|
source = EndianBinaryReader(reader.read(node.size), offset=reader.BaseOffset + node.offset)
|
|
obj = ImportHelper.parse_file(source, parent, name, is_dependency=parent.is_dependency)
|
|
if isinstance(obj, (EndianBinaryReader, SerializedFile)) and parent.environment:
|
|
parent.environment.register_cab(name, obj)
|
|
obj.flags = getattr(node, "flags", 0)
|
|
parent.files[name] = obj
|
|
|
|
BundleFile.read_fs, File.read_files = read_fs, read_files
|
|
try:
|
|
yield
|
|
finally:
|
|
BundleFile.read_fs, File.read_files = old_fs, old_files
|
|
|
|
|
|
def load(path):
|
|
import UnityPy
|
|
with Path(path).open("rb") as stream, streaming_unity():
|
|
return UnityPy.load(stream)
|
|
|
|
|
|
def resource_hashes(bundle):
|
|
from UnityPy.streams import EndianBinaryReader
|
|
result = {}
|
|
for name, obj in bundle.files.items():
|
|
if not isinstance(obj, EndianBinaryReader):
|
|
continue
|
|
before = obj.Position
|
|
obj.Position = 0
|
|
fingerprint = hashlib.sha256()
|
|
while obj.Position < obj.Length:
|
|
fingerprint.update(obj.read(min(1 << 20, obj.Length - obj.Position)))
|
|
obj.Position = before
|
|
result[name] = fingerprint.hexdigest()
|
|
return result
|
|
|
|
|
|
def save(bundle, path):
|
|
from lz4.block import compress
|
|
from UnityPy.streams import EndianBinaryReader, EndianBinaryWriter
|
|
if bundle.signature != "UnityFS" or bundle.parse_version()[:2] != (2019, 4):
|
|
raise ValueError("Unsupported bundle format")
|
|
header = EndianBinaryWriter()
|
|
header.write_string_to_null(bundle.signature)
|
|
header.write_u_int(bundle.version)
|
|
header.write_string_to_null(bundle.version_player)
|
|
header.write_string_to_null(bundle.version_engine)
|
|
size_position = header.Position
|
|
header.write_long(0)
|
|
header.write_u_int(0)
|
|
header.write_u_int(0)
|
|
header.write_u_int(194) # LZ4, directory included, block info at end.
|
|
header.align_stream(16)
|
|
nodes, blocks = [], []
|
|
crc, offset = 0, 0
|
|
path = Path(path)
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
with path.open("w+b") as output, temporary_file() as raw:
|
|
output.write(header.bytes)
|
|
for name, obj in bundle.files.items():
|
|
start = raw.tell()
|
|
if isinstance(obj, EndianBinaryReader):
|
|
obj.Position = 0
|
|
while obj.Position < obj.Length:
|
|
raw.write(obj.read(min(1 << 20, obj.Length - obj.Position)))
|
|
else:
|
|
raw.write(obj.bytes if isinstance(obj, EndianBinaryWriter) else obj.save())
|
|
size = raw.tell() - start
|
|
nodes.append((offset, size, getattr(obj, "flags", 0), name))
|
|
offset += size
|
|
raw.seek(0)
|
|
while data := raw.read(128 << 10):
|
|
packed = compress(data, store_size=False)
|
|
output.write(packed)
|
|
blocks.append((len(data), len(packed), 2))
|
|
crc = zlib.crc32(data, crc)
|
|
info = EndianBinaryWriter(b"\0" * 16)
|
|
info.write_int(len(blocks))
|
|
for raw_size, packed_size, flags in blocks:
|
|
info.write_u_int(raw_size)
|
|
info.write_u_int(packed_size)
|
|
info.write_u_short(flags)
|
|
info.write_int(len(nodes))
|
|
for node_offset, size, flags, name in nodes:
|
|
info.write_long(node_offset)
|
|
info.write_long(size)
|
|
info.write_u_int(flags)
|
|
info.write_string_to_null(name)
|
|
raw_info = info.bytes
|
|
packed_info = compress(raw_info, store_size=False)
|
|
output.write(packed_info)
|
|
size = output.tell()
|
|
output.seek(size_position)
|
|
output.write(struct.pack(">qIII", size, len(packed_info), len(raw_info), 194))
|
|
return crc & 0xffffffff
|