1082 lines
30 KiB
Python
1082 lines
30 KiB
Python
"""Inspect the XOR-wrapped Assimp binary scene caches used by PristonTale EU."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import asdict, dataclass
|
|
import hashlib
|
|
import io
|
|
from pathlib import Path
|
|
import shutil
|
|
import struct
|
|
from typing import BinaryIO, Protocol
|
|
import zlib
|
|
|
|
|
|
ASSBIN_HEADER_LENGTH = 512
|
|
ASSBIN_CHUNK_AITEXTURE = 0x1236
|
|
ASSBIN_CHUNK_AIMESH = 0x1237
|
|
ASSBIN_CHUNK_AISCENE = 0x1239
|
|
ASSBIN_CHUNK_AIANIMATION = 0x123B
|
|
ASSBIN_CHUNK_AINODE = 0x123C
|
|
ASSBIN_CHUNK_AIMATERIAL = 0x123D
|
|
ASSBIN_CHUNK_AIMATERIALPROPERTY = 0x123E
|
|
|
|
PROPERTY_TYPE_NAMES = {
|
|
1: "float",
|
|
2: "double",
|
|
3: "string",
|
|
4: "integer",
|
|
5: "buffer",
|
|
}
|
|
|
|
MAX_STRING_LENGTH = 16 * 1024 * 1024
|
|
MAX_CAPTURE_LENGTH = 16 * 1024 * 1024
|
|
READ_CHUNK_SIZE = 1024 * 1024
|
|
|
|
|
|
class AssbinError(ValueError):
|
|
"""Raised when a wrapped Assbin cache is invalid or unsupported."""
|
|
|
|
|
|
class Reader(Protocol):
|
|
def tell(self) -> int: ...
|
|
|
|
def seek(self, offset: int) -> int: ...
|
|
|
|
def read_exact(self, size: int) -> bytes: ...
|
|
|
|
@property
|
|
def size(self) -> int: ...
|
|
|
|
|
|
class XorFileReader:
|
|
"""Read a file while undoing the position-dependent game wrapper."""
|
|
|
|
def __init__(self, stream: BinaryIO) -> None:
|
|
self._stream = stream
|
|
position = stream.tell()
|
|
stream.seek(0, io.SEEK_END)
|
|
self._size = stream.tell()
|
|
stream.seek(position)
|
|
|
|
@property
|
|
def size(self) -> int:
|
|
return self._size
|
|
|
|
def tell(self) -> int:
|
|
return self._stream.tell()
|
|
|
|
def seek(self, offset: int) -> int:
|
|
if offset < 0 or offset > self._size:
|
|
raise AssbinError(f"seek outside wrapped scene: {offset}")
|
|
return self._stream.seek(offset)
|
|
|
|
def read_exact(self, size: int) -> bytes:
|
|
if size < 0:
|
|
raise AssbinError(f"negative read size: {size}")
|
|
offset = self.tell()
|
|
raw = self._stream.read(size)
|
|
if len(raw) != size:
|
|
raise AssbinError(
|
|
f"unexpected EOF at {offset}: wanted {size} bytes, got {len(raw)}"
|
|
)
|
|
return xor_bytes(raw, offset)
|
|
|
|
|
|
class BytesReader:
|
|
"""Read an already-decoded byte buffer."""
|
|
|
|
def __init__(self, data: bytes) -> None:
|
|
self._stream = io.BytesIO(data)
|
|
self._size = len(data)
|
|
|
|
@property
|
|
def size(self) -> int:
|
|
return self._size
|
|
|
|
def tell(self) -> int:
|
|
return self._stream.tell()
|
|
|
|
def seek(self, offset: int) -> int:
|
|
if offset < 0 or offset > self._size:
|
|
raise AssbinError(f"seek outside decoded scene: {offset}")
|
|
return self._stream.seek(offset)
|
|
|
|
def read_exact(self, size: int) -> bytes:
|
|
offset = self.tell()
|
|
data = self._stream.read(size)
|
|
if len(data) != size:
|
|
raise AssbinError(
|
|
f"unexpected EOF at {offset}: wanted {size} bytes, got {len(data)}"
|
|
)
|
|
return data
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class AssbinHeader:
|
|
signature: str
|
|
version_major: int
|
|
version_minor: int
|
|
version_revision: int
|
|
compile_flags: int
|
|
shortened: bool
|
|
compressed: bool
|
|
source_file: str
|
|
options: str
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Chunk:
|
|
identifier: int
|
|
header_offset: int
|
|
data_offset: int
|
|
data_size: int
|
|
end_offset: int
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class BlobSummary:
|
|
length: int
|
|
sha256: str
|
|
prefix_hex: str
|
|
signature: str | None
|
|
value: object | None
|
|
png: PngMetadata | None
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class PngMetadata:
|
|
valid: bool
|
|
errors: list[str]
|
|
width: int | None
|
|
height: int | None
|
|
bit_depth: int | None
|
|
color_type: int | None
|
|
compression_method: int | None
|
|
filter_method: int | None
|
|
interlace_method: int | None
|
|
chunks: list[str]
|
|
inflated_length: int | None
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class MaterialProperty:
|
|
key: str
|
|
semantic: int
|
|
index: int
|
|
data_length: int
|
|
property_type: int
|
|
property_type_name: str
|
|
data_offset: int
|
|
sha256: str
|
|
prefix_hex: str
|
|
signature: str | None
|
|
value: object | None
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Material:
|
|
index: int
|
|
name: str | None
|
|
properties: list[MaterialProperty]
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class EmbeddedTexture:
|
|
index: int
|
|
chunk_header_offset: int
|
|
chunk_data_size: int
|
|
width: int
|
|
height: int
|
|
format_hint: str
|
|
data_offset: int
|
|
data_length: int
|
|
sha256: str
|
|
prefix_hex: str
|
|
signature: str | None
|
|
png: PngMetadata | None
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class SceneCounts:
|
|
flags: int
|
|
meshes: int
|
|
materials: int
|
|
animations: int
|
|
textures: int
|
|
lights: int
|
|
cameras: int
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class SceneInspection:
|
|
path: str
|
|
file_size: int
|
|
file_sha256: str
|
|
header: AssbinHeader
|
|
counts: SceneCounts
|
|
materials: list[Material]
|
|
embedded_textures: list[EmbeddedTexture]
|
|
scene_chunk_offset: int
|
|
scene_chunk_size: int
|
|
parsed_scene_offset: int
|
|
|
|
def to_dict(self) -> dict[str, object]:
|
|
return asdict(self)
|
|
|
|
|
|
def xor_bytes(data: bytes, absolute_offset: int = 0) -> bytes:
|
|
"""Undo the repeating position-dependent XOR wrapper."""
|
|
|
|
return bytes(
|
|
value ^ ((0x5C + absolute_offset + index) & 0xFF)
|
|
for index, value in enumerate(data)
|
|
)
|
|
|
|
|
|
def inspect_wrapped_assbin(path: Path) -> SceneInspection:
|
|
"""Parse scene metadata and material properties without changing the cache."""
|
|
|
|
resolved_path = path.expanduser().resolve()
|
|
file_sha256 = _hash_path(resolved_path)
|
|
with resolved_path.open("rb") as stream:
|
|
wrapped_reader = XorFileReader(stream)
|
|
header = _read_header(wrapped_reader)
|
|
if header.shortened:
|
|
raise AssbinError("shortened Assbin regression dumps are unsupported")
|
|
|
|
if header.compressed:
|
|
uncompressed_size = _read_u32(wrapped_reader)
|
|
compressed_data = wrapped_reader.read_exact(
|
|
wrapped_reader.size - wrapped_reader.tell()
|
|
)
|
|
try:
|
|
decoded = zlib.decompress(compressed_data)
|
|
except zlib.error as exc:
|
|
raise AssbinError(f"zlib decompression failed: {exc}") from exc
|
|
if len(decoded) != uncompressed_size:
|
|
raise AssbinError(
|
|
"decompressed size mismatch: "
|
|
f"expected {uncompressed_size}, got {len(decoded)}"
|
|
)
|
|
reader: Reader = BytesReader(decoded)
|
|
else:
|
|
reader = wrapped_reader
|
|
|
|
counts, materials, textures, scene_chunk, parsed_offset = _read_scene(reader)
|
|
return SceneInspection(
|
|
path=str(resolved_path),
|
|
file_size=wrapped_reader.size,
|
|
file_sha256=file_sha256,
|
|
header=header,
|
|
counts=counts,
|
|
materials=materials,
|
|
embedded_textures=textures,
|
|
scene_chunk_offset=scene_chunk.header_offset,
|
|
scene_chunk_size=scene_chunk.data_size,
|
|
parsed_scene_offset=parsed_offset,
|
|
)
|
|
|
|
|
|
def extract_embedded_texture(
|
|
path: Path,
|
|
inspection: SceneInspection,
|
|
texture_index: int,
|
|
output: Path,
|
|
) -> dict[str, object]:
|
|
"""Extract one decoded embedded texture from an uncompressed wrapped cache."""
|
|
|
|
if inspection.header.compressed:
|
|
raise AssbinError("direct extraction from compressed Assbin caches is unsupported")
|
|
if texture_index < 0 or texture_index >= len(inspection.embedded_textures):
|
|
raise AssbinError(
|
|
f"texture index {texture_index} is outside 0.."
|
|
f"{len(inspection.embedded_textures) - 1}"
|
|
)
|
|
texture = inspection.embedded_textures[texture_index]
|
|
resolved_path = path.expanduser().resolve()
|
|
resolved_output = output.expanduser().resolve()
|
|
resolved_output.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
digest = hashlib.sha256()
|
|
written = 0
|
|
try:
|
|
with resolved_path.open("rb") as source, resolved_output.open("xb") as target:
|
|
reader = XorFileReader(source)
|
|
reader.seek(texture.data_offset)
|
|
remaining = texture.data_length
|
|
while remaining:
|
|
part = reader.read_exact(min(remaining, READ_CHUNK_SIZE))
|
|
target.write(part)
|
|
digest.update(part)
|
|
written += len(part)
|
|
remaining -= len(part)
|
|
except BaseException:
|
|
resolved_output.unlink(missing_ok=True)
|
|
raise
|
|
|
|
actual_sha256 = digest.hexdigest()
|
|
if written != texture.data_length or actual_sha256 != texture.sha256:
|
|
resolved_output.unlink(missing_ok=True)
|
|
raise AssbinError(
|
|
"extracted texture verification failed: "
|
|
f"size {written}/{texture.data_length}, "
|
|
f"sha256 {actual_sha256}/{texture.sha256}"
|
|
)
|
|
return {
|
|
"index": texture_index,
|
|
"output": str(resolved_output),
|
|
"size": written,
|
|
"sha256": actual_sha256,
|
|
"format_hint": texture.format_hint,
|
|
"signature": texture.signature,
|
|
}
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class _BinaryPatch:
|
|
offset: int
|
|
old_length: int
|
|
data: bytes
|
|
description: str
|
|
|
|
|
|
def normalize_wrapped_assbin_textures(
|
|
source: Path,
|
|
output: Path,
|
|
texture_indices: tuple[int, ...],
|
|
) -> dict[str, object]:
|
|
"""Convert selected embedded 16-bit RGB PNGs to 8-bit in a new cache."""
|
|
|
|
if not texture_indices:
|
|
raise AssbinError("at least one texture index is required")
|
|
if len(set(texture_indices)) != len(texture_indices):
|
|
raise AssbinError("texture indices must be unique")
|
|
resolved_source = source.expanduser().resolve(strict=True)
|
|
resolved_output = output.expanduser().resolve()
|
|
if resolved_source == resolved_output:
|
|
raise AssbinError("normalization output must differ from its source")
|
|
if resolved_output.exists() or resolved_output.is_symlink():
|
|
raise AssbinError(f"normalization output already exists: {resolved_output}")
|
|
|
|
inspection = inspect_wrapped_assbin(resolved_source)
|
|
if inspection.header.compressed:
|
|
raise AssbinError("normalizing compressed Assbin caches is unsupported")
|
|
selected = []
|
|
for index in sorted(texture_indices):
|
|
if index < 0 or index >= len(inspection.embedded_textures):
|
|
raise AssbinError(
|
|
f"texture index {index} is outside 0.."
|
|
f"{len(inspection.embedded_textures) - 1}"
|
|
)
|
|
selected.append(inspection.embedded_textures[index])
|
|
|
|
patches: list[_BinaryPatch] = []
|
|
changes: list[dict[str, object]] = []
|
|
total_delta = 0
|
|
with resolved_source.open("rb") as stream:
|
|
reader = XorFileReader(stream)
|
|
for texture in selected:
|
|
reader.seek(texture.data_offset)
|
|
old_png = reader.read_exact(texture.data_length)
|
|
new_png = normalize_png_16_rgb_to_8(old_png)
|
|
new_metadata = _inspect_png(new_png)
|
|
new_sha256 = hashlib.sha256(new_png).hexdigest()
|
|
delta = len(new_png) - texture.data_length
|
|
total_delta += delta
|
|
patches.extend(
|
|
(
|
|
_BinaryPatch(
|
|
texture.chunk_header_offset + 4,
|
|
4,
|
|
struct.pack("<I", texture.chunk_data_size + delta),
|
|
f"texture {texture.index} chunk size",
|
|
),
|
|
_BinaryPatch(
|
|
texture.data_offset - 16,
|
|
4,
|
|
struct.pack("<I", len(new_png)),
|
|
f"texture {texture.index} compressed byte count",
|
|
),
|
|
_BinaryPatch(
|
|
texture.data_offset,
|
|
texture.data_length,
|
|
new_png,
|
|
f"texture {texture.index} PNG data",
|
|
),
|
|
)
|
|
)
|
|
changes.append(
|
|
{
|
|
"index": texture.index,
|
|
"old_size": texture.data_length,
|
|
"new_size": len(new_png),
|
|
"size_delta": delta,
|
|
"old_sha256": texture.sha256,
|
|
"new_sha256": new_sha256,
|
|
"new_scanlines_sha256": _png_scanlines_sha256(new_png),
|
|
"old_png": asdict(texture.png) if texture.png is not None else None,
|
|
"new_png": asdict(new_metadata),
|
|
}
|
|
)
|
|
|
|
patches.append(
|
|
_BinaryPatch(
|
|
inspection.scene_chunk_offset + 4,
|
|
4,
|
|
struct.pack("<I", inspection.scene_chunk_size + total_delta),
|
|
"scene chunk size",
|
|
)
|
|
)
|
|
resolved_output.parent.mkdir(parents=True, exist_ok=True)
|
|
try:
|
|
output_sha256, output_size = _rewrite_wrapped_file(
|
|
resolved_source,
|
|
resolved_output,
|
|
patches,
|
|
)
|
|
shutil.copymode(resolved_source, resolved_output)
|
|
verification = inspect_wrapped_assbin(resolved_output)
|
|
_verify_normalization(inspection, verification, set(texture_indices))
|
|
except BaseException:
|
|
resolved_output.unlink(missing_ok=True)
|
|
raise
|
|
|
|
return {
|
|
"source": str(resolved_source),
|
|
"output": str(resolved_output),
|
|
"source_size": inspection.file_size,
|
|
"output_size": output_size,
|
|
"size_delta": total_delta,
|
|
"source_sha256": inspection.file_sha256,
|
|
"output_sha256": output_sha256,
|
|
"source_scene_chunk_size": inspection.scene_chunk_size,
|
|
"output_scene_chunk_size": verification.scene_chunk_size,
|
|
"texture_indices": list(texture_indices),
|
|
"changes": changes,
|
|
"verified": True,
|
|
}
|
|
|
|
|
|
def normalize_png_16_rgb_to_8(data: bytes) -> bytes:
|
|
"""Losslessly parse a 16-bit RGB PNG and downsample its channels to 8-bit."""
|
|
|
|
metadata = _inspect_png(data)
|
|
if not metadata.valid:
|
|
raise AssbinError(f"cannot normalize invalid PNG: {'; '.join(metadata.errors)}")
|
|
if (
|
|
metadata.width is None
|
|
or metadata.height is None
|
|
or metadata.bit_depth != 16
|
|
or metadata.color_type != 2
|
|
or metadata.interlace_method != 0
|
|
):
|
|
raise AssbinError(
|
|
"normalization requires a non-interlaced 16-bit RGB PNG; got "
|
|
f"{metadata.width}x{metadata.height}, bit depth {metadata.bit_depth}, "
|
|
f"color type {metadata.color_type}, interlace {metadata.interlace_method}"
|
|
)
|
|
|
|
chunks = _png_chunks(data)
|
|
compressed = b"".join(chunk_data for name, chunk_data in chunks if name == b"IDAT")
|
|
try:
|
|
filtered = zlib.decompress(compressed)
|
|
except zlib.error as exc:
|
|
raise AssbinError(f"PNG IDAT decompression failed: {exc}") from exc
|
|
row_bytes = metadata.width * 6
|
|
expected_length = (row_bytes + 1) * metadata.height
|
|
if len(filtered) != expected_length:
|
|
raise AssbinError(
|
|
f"unexpected 16-bit RGB scanline length: {len(filtered)} != {expected_length}"
|
|
)
|
|
|
|
converted = bytearray()
|
|
previous = bytes(row_bytes)
|
|
offset = 0
|
|
for _ in range(metadata.height):
|
|
filter_type = filtered[offset]
|
|
row = bytearray(filtered[offset + 1 : offset + 1 + row_bytes])
|
|
_unfilter_png_row(row, previous, filter_type, 6)
|
|
converted.append(0)
|
|
for byte_index in range(0, row_bytes, 2):
|
|
sample = (row[byte_index] << 8) | row[byte_index + 1]
|
|
converted.append((sample + 128) // 257)
|
|
previous = bytes(row)
|
|
offset += row_bytes + 1
|
|
|
|
new_idat = zlib.compress(bytes(converted), level=9)
|
|
result = bytearray(b"\x89PNG\r\n\x1a\n")
|
|
wrote_idat = False
|
|
for name, chunk_data in chunks:
|
|
if name == b"IHDR":
|
|
mutable_header = bytearray(chunk_data)
|
|
mutable_header[8] = 8
|
|
result.extend(_make_png_chunk(name, bytes(mutable_header)))
|
|
elif name == b"IDAT":
|
|
if not wrote_idat:
|
|
result.extend(_make_png_chunk(b"IDAT", new_idat))
|
|
wrote_idat = True
|
|
else:
|
|
result.extend(_make_png_chunk(name, chunk_data))
|
|
normalized = bytes(result)
|
|
new_metadata = _inspect_png(normalized)
|
|
if not new_metadata.valid or new_metadata.bit_depth != 8:
|
|
raise AssbinError(
|
|
f"normalized PNG failed verification: {'; '.join(new_metadata.errors)}"
|
|
)
|
|
return normalized
|
|
|
|
|
|
def _read_header(reader: Reader) -> AssbinHeader:
|
|
if reader.size < ASSBIN_HEADER_LENGTH:
|
|
raise AssbinError(
|
|
f"wrapped scene is smaller than the {ASSBIN_HEADER_LENGTH}-byte header"
|
|
)
|
|
reader.seek(0)
|
|
signature_raw = reader.read_exact(44)
|
|
if not signature_raw.startswith(b"ASSIMP.binary-dump."):
|
|
raise AssbinError(
|
|
f"decoded header has no Assbin signature: {signature_raw[:20]!r}"
|
|
)
|
|
version_major = _read_u32(reader)
|
|
version_minor = _read_u32(reader)
|
|
version_revision = _read_u32(reader)
|
|
compile_flags = _read_u32(reader)
|
|
shortened = _read_u16(reader) > 0
|
|
compressed = _read_u16(reader) > 0
|
|
source_file = _decode_fixed(reader.read_exact(256))
|
|
options = _decode_fixed(reader.read_exact(128))
|
|
reader.read_exact(64)
|
|
if reader.tell() != ASSBIN_HEADER_LENGTH:
|
|
raise AssbinError(f"internal header-size error: {reader.tell()}")
|
|
return AssbinHeader(
|
|
signature=signature_raw.rstrip(b"\x00\r\n").decode("utf-8", "replace"),
|
|
version_major=version_major,
|
|
version_minor=version_minor,
|
|
version_revision=version_revision,
|
|
compile_flags=compile_flags,
|
|
shortened=shortened,
|
|
compressed=compressed,
|
|
source_file=source_file,
|
|
options=options,
|
|
)
|
|
|
|
|
|
def _read_scene(
|
|
reader: Reader,
|
|
) -> tuple[SceneCounts, list[Material], list[EmbeddedTexture], Chunk, int]:
|
|
scene_chunk = _read_chunk(reader, ASSBIN_CHUNK_AISCENE)
|
|
counts = SceneCounts(
|
|
flags=_read_u32(reader),
|
|
meshes=_read_u32(reader),
|
|
materials=_read_u32(reader),
|
|
animations=_read_u32(reader),
|
|
textures=_read_u32(reader),
|
|
lights=_read_u32(reader),
|
|
cameras=_read_u32(reader),
|
|
)
|
|
|
|
_skip_chunk(reader, ASSBIN_CHUNK_AINODE, scene_chunk.end_offset)
|
|
for _ in range(counts.meshes):
|
|
_skip_chunk(reader, ASSBIN_CHUNK_AIMESH, scene_chunk.end_offset)
|
|
|
|
materials = [
|
|
_read_material(reader, index, scene_chunk.end_offset)
|
|
for index in range(counts.materials)
|
|
]
|
|
|
|
for _ in range(counts.animations):
|
|
_skip_chunk(reader, ASSBIN_CHUNK_AIANIMATION, scene_chunk.end_offset)
|
|
|
|
textures = [
|
|
_read_texture(reader, index, scene_chunk.end_offset)
|
|
for index in range(counts.textures)
|
|
]
|
|
|
|
# Light and camera layouts are irrelevant to texture diagnosis. Their complete
|
|
# chunk sizes let us advance safely without depending on their struct versions.
|
|
for _ in range(counts.lights + counts.cameras):
|
|
_skip_unknown_chunk(reader, scene_chunk.end_offset)
|
|
|
|
parsed_offset = reader.tell()
|
|
if parsed_offset > scene_chunk.end_offset:
|
|
raise AssbinError(
|
|
f"scene parser crossed chunk boundary: {parsed_offset} > {scene_chunk.end_offset}"
|
|
)
|
|
return counts, materials, textures, scene_chunk, parsed_offset
|
|
|
|
|
|
def _read_material(reader: Reader, index: int, parent_end: int) -> Material:
|
|
chunk = _read_chunk(reader, ASSBIN_CHUNK_AIMATERIAL, parent_end)
|
|
property_count = _read_u32(reader)
|
|
properties = [
|
|
_read_material_property(reader, chunk.end_offset)
|
|
for _ in range(property_count)
|
|
]
|
|
_assert_not_past(reader, chunk)
|
|
reader.seek(chunk.end_offset)
|
|
name = _material_name(properties)
|
|
return Material(index=index, name=name, properties=properties)
|
|
|
|
|
|
def _read_material_property(reader: Reader, parent_end: int) -> MaterialProperty:
|
|
chunk = _read_chunk(reader, ASSBIN_CHUNK_AIMATERIALPROPERTY, parent_end)
|
|
key = _read_string(reader)
|
|
semantic = _read_u32(reader)
|
|
index = _read_u32(reader)
|
|
data_length = _read_u32(reader)
|
|
property_type = _read_u32(reader)
|
|
data_offset = reader.tell()
|
|
if data_offset + data_length > chunk.end_offset:
|
|
raise AssbinError(
|
|
f"material property {key!r} data exceeds its chunk: "
|
|
f"{data_offset} + {data_length} > {chunk.end_offset}"
|
|
)
|
|
blob = _summarize_blob(reader, data_length, property_type)
|
|
_assert_not_past(reader, chunk)
|
|
reader.seek(chunk.end_offset)
|
|
return MaterialProperty(
|
|
key=key,
|
|
semantic=semantic,
|
|
index=index,
|
|
data_length=data_length,
|
|
property_type=property_type,
|
|
property_type_name=PROPERTY_TYPE_NAMES.get(property_type, "unknown"),
|
|
data_offset=data_offset,
|
|
sha256=blob.sha256,
|
|
prefix_hex=blob.prefix_hex,
|
|
signature=blob.signature,
|
|
value=blob.value,
|
|
)
|
|
|
|
|
|
def _read_texture(reader: Reader, index: int, parent_end: int) -> EmbeddedTexture:
|
|
chunk = _read_chunk(reader, ASSBIN_CHUNK_AITEXTURE, parent_end)
|
|
width = _read_u32(reader)
|
|
height = _read_u32(reader)
|
|
format_hint = _decode_fixed(reader.read_exact(8))
|
|
data_length = width if height == 0 else width * height * 4
|
|
data_offset = reader.tell()
|
|
if data_offset + data_length > chunk.end_offset:
|
|
raise AssbinError(
|
|
f"embedded texture {index} data exceeds its chunk boundary"
|
|
)
|
|
blob = _summarize_blob(reader, data_length, 5)
|
|
_assert_not_past(reader, chunk)
|
|
reader.seek(chunk.end_offset)
|
|
return EmbeddedTexture(
|
|
index=index,
|
|
chunk_header_offset=chunk.header_offset,
|
|
chunk_data_size=chunk.data_size,
|
|
width=width,
|
|
height=height,
|
|
format_hint=format_hint,
|
|
data_offset=data_offset,
|
|
data_length=data_length,
|
|
sha256=blob.sha256,
|
|
prefix_hex=blob.prefix_hex,
|
|
signature=blob.signature,
|
|
png=blob.png,
|
|
)
|
|
|
|
|
|
def _summarize_blob(reader: Reader, length: int, property_type: int) -> BlobSummary:
|
|
remaining = length
|
|
digest = hashlib.sha256()
|
|
prefix = bytearray()
|
|
captured = bytearray()
|
|
while remaining:
|
|
part = reader.read_exact(min(remaining, READ_CHUNK_SIZE))
|
|
digest.update(part)
|
|
if len(prefix) < 64:
|
|
prefix.extend(part[: 64 - len(prefix)])
|
|
if length <= MAX_CAPTURE_LENGTH:
|
|
captured.extend(part)
|
|
remaining -= len(part)
|
|
data = bytes(captured) if length <= MAX_CAPTURE_LENGTH else bytes(prefix)
|
|
signature = _identify_signature(bytes(prefix))
|
|
return BlobSummary(
|
|
length=length,
|
|
sha256=digest.hexdigest(),
|
|
prefix_hex=bytes(prefix).hex(),
|
|
signature=signature,
|
|
value=_decode_property_value(data, length, property_type),
|
|
png=_inspect_png(data) if signature == "PNG" and len(data) == length else None,
|
|
)
|
|
|
|
|
|
def _decode_property_value(data: bytes, full_length: int, property_type: int) -> object | None:
|
|
if full_length > MAX_CAPTURE_LENGTH:
|
|
return None
|
|
if property_type == 3 and len(data) >= 4:
|
|
string_length = struct.unpack_from("<I", data)[0]
|
|
if string_length <= len(data) - 4:
|
|
return data[4 : 4 + string_length].rstrip(b"\x00").decode(
|
|
"utf-8", "replace"
|
|
)
|
|
if property_type == 4 and len(data) % 4 == 0:
|
|
values = list(struct.unpack(f"<{len(data) // 4}I", data))
|
|
return values[0] if len(values) == 1 else values
|
|
if property_type == 1 and len(data) % 4 == 0:
|
|
values = list(struct.unpack(f"<{len(data) // 4}f", data))
|
|
return values[0] if len(values) == 1 else values
|
|
if property_type == 2 and len(data) % 8 == 0:
|
|
values = list(struct.unpack(f"<{len(data) // 8}d", data))
|
|
return values[0] if len(values) == 1 else values
|
|
return None
|
|
|
|
|
|
def _identify_signature(prefix: bytes) -> str | None:
|
|
signatures = (
|
|
(b"\x89PNG\r\n\x1a\n", "PNG"),
|
|
(b"\xff\xd8\xff", "JPEG"),
|
|
(b"BM", "BMP"),
|
|
(b"DDS ", "DDS"),
|
|
(b"GIF87a", "GIF87a"),
|
|
(b"GIF89a", "GIF89a"),
|
|
(b"RIFF", "RIFF"),
|
|
(b"KTX 11\xbb\r\n\x1a\n", "KTX"),
|
|
)
|
|
for magic, name in signatures:
|
|
if prefix.startswith(magic):
|
|
return name
|
|
return None
|
|
|
|
|
|
def _inspect_png(data: bytes) -> PngMetadata:
|
|
errors: list[str] = []
|
|
chunks: list[str] = []
|
|
width: int | None = None
|
|
height: int | None = None
|
|
bit_depth: int | None = None
|
|
color_type: int | None = None
|
|
compression_method: int | None = None
|
|
filter_method: int | None = None
|
|
interlace_method: int | None = None
|
|
idat_parts: list[bytes] = []
|
|
inflated_length: int | None = None
|
|
offset = 8
|
|
saw_iend = False
|
|
|
|
if not data.startswith(b"\x89PNG\r\n\x1a\n"):
|
|
errors.append("invalid PNG signature")
|
|
else:
|
|
while offset < len(data):
|
|
if offset + 12 > len(data):
|
|
errors.append(f"truncated PNG chunk header at {offset}")
|
|
break
|
|
chunk_length = struct.unpack_from(">I", data, offset)[0]
|
|
chunk_type = data[offset + 4 : offset + 8]
|
|
data_start = offset + 8
|
|
data_end = data_start + chunk_length
|
|
chunk_end = data_end + 4
|
|
if chunk_end > len(data):
|
|
errors.append(
|
|
f"PNG chunk {chunk_type!r} at {offset} exceeds data length"
|
|
)
|
|
break
|
|
chunk_data = data[data_start:data_end]
|
|
stored_crc = struct.unpack_from(">I", data, data_end)[0]
|
|
calculated_crc = zlib.crc32(chunk_type)
|
|
calculated_crc = zlib.crc32(chunk_data, calculated_crc) & 0xFFFFFFFF
|
|
name = chunk_type.decode("ascii", "replace")
|
|
chunks.append(name)
|
|
if stored_crc != calculated_crc:
|
|
errors.append(
|
|
f"CRC mismatch in {name} chunk at {offset}: "
|
|
f"0x{stored_crc:08x} != 0x{calculated_crc:08x}"
|
|
)
|
|
|
|
if chunk_type == b"IHDR":
|
|
if chunk_length != 13:
|
|
errors.append(f"IHDR length is {chunk_length}, expected 13")
|
|
else:
|
|
(
|
|
width,
|
|
height,
|
|
bit_depth,
|
|
color_type,
|
|
compression_method,
|
|
filter_method,
|
|
interlace_method,
|
|
) = struct.unpack(">IIBBBBB", chunk_data)
|
|
elif chunk_type == b"IDAT":
|
|
idat_parts.append(chunk_data)
|
|
elif chunk_type == b"IEND":
|
|
saw_iend = True
|
|
if chunk_length != 0:
|
|
errors.append("IEND chunk is not empty")
|
|
offset = chunk_end
|
|
break
|
|
offset = chunk_end
|
|
|
|
if not chunks or chunks[0] != "IHDR":
|
|
errors.append("IHDR is missing or is not the first chunk")
|
|
if not idat_parts:
|
|
errors.append("IDAT is missing")
|
|
else:
|
|
try:
|
|
inflated_length = len(zlib.decompress(b"".join(idat_parts)))
|
|
except zlib.error as exc:
|
|
errors.append(f"IDAT zlib stream is invalid: {exc}")
|
|
if not saw_iend:
|
|
errors.append("IEND is missing")
|
|
elif offset != len(data):
|
|
errors.append(f"{len(data) - offset} trailing bytes after IEND")
|
|
|
|
return PngMetadata(
|
|
valid=not errors,
|
|
errors=errors,
|
|
width=width,
|
|
height=height,
|
|
bit_depth=bit_depth,
|
|
color_type=color_type,
|
|
compression_method=compression_method,
|
|
filter_method=filter_method,
|
|
interlace_method=interlace_method,
|
|
chunks=chunks,
|
|
inflated_length=inflated_length,
|
|
)
|
|
|
|
|
|
def _png_chunks(data: bytes) -> list[tuple[bytes, bytes]]:
|
|
chunks: list[tuple[bytes, bytes]] = []
|
|
offset = 8
|
|
while offset < len(data):
|
|
if offset + 12 > len(data):
|
|
raise AssbinError(f"truncated PNG chunk at {offset}")
|
|
length = struct.unpack_from(">I", data, offset)[0]
|
|
name = data[offset + 4 : offset + 8]
|
|
data_start = offset + 8
|
|
data_end = data_start + length
|
|
chunk_end = data_end + 4
|
|
if chunk_end > len(data):
|
|
raise AssbinError(f"PNG chunk {name!r} exceeds data length")
|
|
chunks.append((name, data[data_start:data_end]))
|
|
offset = chunk_end
|
|
if name == b"IEND":
|
|
break
|
|
if offset != len(data):
|
|
raise AssbinError(f"PNG has {len(data) - offset} trailing bytes")
|
|
return chunks
|
|
|
|
|
|
def _png_scanlines_sha256(data: bytes) -> str:
|
|
"""Hash decoded PNG scanlines independently of the zlib encoder bytes."""
|
|
|
|
compressed = b"".join(
|
|
chunk_data
|
|
for name, chunk_data in _png_chunks(data)
|
|
if name == b"IDAT"
|
|
)
|
|
try:
|
|
scanlines = zlib.decompress(compressed)
|
|
except zlib.error as exc:
|
|
raise AssbinError(f"PNG IDAT decompression failed: {exc}") from exc
|
|
return hashlib.sha256(scanlines).hexdigest()
|
|
|
|
|
|
def _make_png_chunk(name: bytes, data: bytes) -> bytes:
|
|
crc = zlib.crc32(name)
|
|
crc = zlib.crc32(data, crc) & 0xFFFFFFFF
|
|
return struct.pack(">I", len(data)) + name + data + struct.pack(">I", crc)
|
|
|
|
|
|
def _unfilter_png_row(
|
|
row: bytearray,
|
|
previous: bytes,
|
|
filter_type: int,
|
|
bytes_per_pixel: int,
|
|
) -> None:
|
|
if filter_type == 0:
|
|
return
|
|
if filter_type not in {1, 2, 3, 4}:
|
|
raise AssbinError(f"unsupported PNG filter type: {filter_type}")
|
|
for index in range(len(row)):
|
|
left = row[index - bytes_per_pixel] if index >= bytes_per_pixel else 0
|
|
up = previous[index]
|
|
upper_left = previous[index - bytes_per_pixel] if index >= bytes_per_pixel else 0
|
|
if filter_type == 1:
|
|
predictor = left
|
|
elif filter_type == 2:
|
|
predictor = up
|
|
elif filter_type == 3:
|
|
predictor = (left + up) // 2
|
|
else:
|
|
predictor = _paeth_predictor(left, up, upper_left)
|
|
row[index] = (row[index] + predictor) & 0xFF
|
|
|
|
|
|
def _paeth_predictor(left: int, up: int, upper_left: int) -> int:
|
|
prediction = left + up - upper_left
|
|
left_distance = abs(prediction - left)
|
|
up_distance = abs(prediction - up)
|
|
upper_left_distance = abs(prediction - upper_left)
|
|
if left_distance <= up_distance and left_distance <= upper_left_distance:
|
|
return left
|
|
if up_distance <= upper_left_distance:
|
|
return up
|
|
return upper_left
|
|
|
|
|
|
def _rewrite_wrapped_file(
|
|
source: Path,
|
|
output: Path,
|
|
patches: list[_BinaryPatch],
|
|
) -> tuple[str, int]:
|
|
ordered = sorted(patches, key=lambda item: item.offset)
|
|
previous_end = 0
|
|
for patch in ordered:
|
|
if patch.offset < previous_end:
|
|
raise AssbinError(
|
|
f"overlapping binary patch at {patch.offset}: {patch.description}"
|
|
)
|
|
previous_end = patch.offset + patch.old_length
|
|
|
|
digest = hashlib.sha256()
|
|
output_offset = 0
|
|
try:
|
|
with source.open("rb") as source_stream, output.open("xb") as output_stream:
|
|
reader = XorFileReader(source_stream)
|
|
for patch in ordered:
|
|
copy_length = patch.offset - reader.tell()
|
|
output_offset = _copy_rewrapped(
|
|
reader,
|
|
output_stream,
|
|
digest,
|
|
copy_length,
|
|
output_offset,
|
|
)
|
|
reader.seek(patch.offset + patch.old_length)
|
|
wrapped_patch = xor_bytes(patch.data, output_offset)
|
|
output_stream.write(wrapped_patch)
|
|
digest.update(wrapped_patch)
|
|
output_offset += len(patch.data)
|
|
output_offset = _copy_rewrapped(
|
|
reader,
|
|
output_stream,
|
|
digest,
|
|
reader.size - reader.tell(),
|
|
output_offset,
|
|
)
|
|
except BaseException:
|
|
output.unlink(missing_ok=True)
|
|
raise
|
|
return digest.hexdigest(), output_offset
|
|
|
|
|
|
def _copy_rewrapped(
|
|
reader: Reader,
|
|
output: BinaryIO,
|
|
digest: object,
|
|
length: int,
|
|
output_offset: int,
|
|
) -> int:
|
|
remaining = length
|
|
while remaining:
|
|
decoded = reader.read_exact(min(remaining, READ_CHUNK_SIZE))
|
|
wrapped = xor_bytes(decoded, output_offset)
|
|
output.write(wrapped)
|
|
digest.update(wrapped) # type: ignore[attr-defined]
|
|
output_offset += len(wrapped)
|
|
remaining -= len(decoded)
|
|
return output_offset
|
|
|
|
|
|
def _verify_normalization(
|
|
source: SceneInspection,
|
|
output: SceneInspection,
|
|
selected_indices: set[int],
|
|
) -> None:
|
|
if source.counts != output.counts:
|
|
raise AssbinError("normalized scene counts differ from source")
|
|
if len(source.embedded_textures) != len(output.embedded_textures):
|
|
raise AssbinError("normalized embedded texture count differs from source")
|
|
for before, after in zip(source.embedded_textures, output.embedded_textures):
|
|
if before.index in selected_indices:
|
|
if after.png is None or not after.png.valid or after.png.bit_depth != 8:
|
|
raise AssbinError(
|
|
f"normalized texture {before.index} is not a valid 8-bit PNG"
|
|
)
|
|
elif before.sha256 != after.sha256 or before.data_length != after.data_length:
|
|
raise AssbinError(f"unselected texture {before.index} changed")
|
|
|
|
|
|
def _material_name(properties: list[MaterialProperty]) -> str | None:
|
|
for prop in properties:
|
|
if prop.key in {"?mat.name", "$mat.name"} and isinstance(prop.value, str):
|
|
return prop.value
|
|
return None
|
|
|
|
|
|
def _read_chunk(reader: Reader, expected: int, parent_end: int | None = None) -> Chunk:
|
|
header_offset = reader.tell()
|
|
identifier = _read_u32(reader)
|
|
data_size = _read_u32(reader)
|
|
data_offset = reader.tell()
|
|
end_offset = data_offset + data_size
|
|
limit = reader.size if parent_end is None else min(reader.size, parent_end)
|
|
if identifier != expected:
|
|
raise AssbinError(
|
|
f"unexpected chunk at {header_offset}: expected 0x{expected:04x}, "
|
|
f"got 0x{identifier:04x}"
|
|
)
|
|
if end_offset > limit:
|
|
raise AssbinError(
|
|
f"chunk 0x{identifier:04x} at {header_offset} ends outside parent: "
|
|
f"{end_offset} > {limit}"
|
|
)
|
|
return Chunk(identifier, header_offset, data_offset, data_size, end_offset)
|
|
|
|
|
|
def _skip_chunk(reader: Reader, expected: int, parent_end: int) -> None:
|
|
chunk = _read_chunk(reader, expected, parent_end)
|
|
reader.seek(chunk.end_offset)
|
|
|
|
|
|
def _skip_unknown_chunk(reader: Reader, parent_end: int) -> None:
|
|
header_offset = reader.tell()
|
|
identifier = _read_u32(reader)
|
|
data_size = _read_u32(reader)
|
|
end_offset = reader.tell() + data_size
|
|
if end_offset > parent_end or end_offset > reader.size:
|
|
raise AssbinError(
|
|
f"chunk 0x{identifier:04x} at {header_offset} ends outside scene"
|
|
)
|
|
reader.seek(end_offset)
|
|
|
|
|
|
def _assert_not_past(reader: Reader, chunk: Chunk) -> None:
|
|
if reader.tell() > chunk.end_offset:
|
|
raise AssbinError(
|
|
f"parser crossed chunk 0x{chunk.identifier:04x} boundary: "
|
|
f"{reader.tell()} > {chunk.end_offset}"
|
|
)
|
|
|
|
|
|
def _read_string(reader: Reader) -> str:
|
|
length = _read_u32(reader)
|
|
if length > MAX_STRING_LENGTH:
|
|
raise AssbinError(f"Assbin string is unreasonably large: {length}")
|
|
return reader.read_exact(length).decode("utf-8", "replace")
|
|
|
|
|
|
def _read_u16(reader: Reader) -> int:
|
|
return struct.unpack("<H", reader.read_exact(2))[0]
|
|
|
|
|
|
def _read_u32(reader: Reader) -> int:
|
|
return struct.unpack("<I", reader.read_exact(4))[0]
|
|
|
|
|
|
def _decode_fixed(data: bytes) -> str:
|
|
return data.split(b"\x00", 1)[0].decode("utf-8", "replace")
|
|
|
|
|
|
def _hash_path(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as stream:
|
|
while chunk := stream.read(READ_CHUNK_SIZE):
|
|
digest.update(chunk)
|
|
return digest.hexdigest()
|