111 lines
4.2 KiB
Python
111 lines
4.2 KiB
Python
"""LZO, UE3 packages and soundbanks (ackit/lzo.py, upk.py, bnk.py)."""
|
|
import os
|
|
|
|
import pytest
|
|
|
|
import synth
|
|
from ackit import lzo
|
|
from ackit.bnk import Bank, find_banks, fnv1, IN_MEMORY
|
|
from ackit.upk import Package
|
|
|
|
|
|
# --- LZO ----------------------------------------------------------------------------------------
|
|
@pytest.mark.parametrize("n", [0, 1, 3, 4, 17, 18, 19, 238, 239, 273, 274, 529, 0x20000])
|
|
def test_lzo_stored_roundtrip(n):
|
|
data = os.urandom(n)
|
|
assert lzo.decompress(lzo.compress_stored(data), n) == data
|
|
|
|
|
|
def test_lzo_matches():
|
|
# "abc" + M3 match (9 bytes at distance 3) + end
|
|
assert lzo.decompress(bytes([20]) + b"abc" + bytes([39, 0x08, 0x00]) + lzo.EOF) == b"abc" * 4
|
|
# "abcd" + M2 match (4 bytes at distance 4) + end
|
|
assert lzo.decompress(bytes([21]) + b"abcd" + bytes([108, 0]) + lzo.EOF) == b"abcd" * 2
|
|
|
|
|
|
def test_lzo_errors():
|
|
with pytest.raises(lzo.LzoError):
|
|
lzo.decompress(bytes([25]) + b"abc") # truncated
|
|
with pytest.raises(lzo.LzoError):
|
|
lzo.decompress(lzo.compress_stored(b"hello") + b"x") # trailing bytes
|
|
with pytest.raises(lzo.LzoError):
|
|
lzo.decompress(lzo.compress_stored(b"hello"), 6) # wrong size
|
|
|
|
|
|
# --- packages -----------------------------------------------------------------------------------
|
|
@pytest.fixture
|
|
def pkg(tmp_path):
|
|
body = bytes(range(256)) * 80 # 20480 bytes: 5 blocks, 2 chunks
|
|
p = str(tmp_path / "Synth.upk")
|
|
start = synth.package(p, body)
|
|
return p, body, start
|
|
|
|
|
|
def test_package_read(pkg):
|
|
p, body, start = pkg
|
|
pk = Package(p)
|
|
assert pk.compressed and pk.version == 805 and pk.header_size == start
|
|
assert pk.size == start + len(body) and len(pk.chunks) == 2 and len(pk.blocks) == 5
|
|
assert pk.read(start + 4000, 300) == body[4000:4300] # crosses a block boundary
|
|
assert pk.read_all() == (start, body)
|
|
|
|
|
|
def test_package_patch(pkg):
|
|
p, body, start = pkg
|
|
raw = open(p, "rb").read()
|
|
pk = Package(p)
|
|
assert pk.patched([]) == raw # nothing changed -> same bytes
|
|
new = pk.patched([(start + 16334, b"X" * 100)]) # second chunk, across two blocks
|
|
q = p + ".new"
|
|
open(q, "wb").write(new)
|
|
pk2 = Package(q)
|
|
exp = bytearray(body)
|
|
exp[16334:16434] = b"X" * 100
|
|
assert pk2.read_all()[1] == bytes(exp)
|
|
# the first chunk is copied verbatim
|
|
c0 = pk.chunks[0]
|
|
assert new[c0[2]:c0[2] + c0[3]] == raw[c0[2]:c0[2] + c0[3]]
|
|
with pytest.raises(ValueError):
|
|
pk.patched([(start + len(body) - 10, b"too long" * 4)])
|
|
|
|
|
|
# --- banks --------------------------------------------------------------------------------------
|
|
def test_bank_parse():
|
|
media = {0x111: b"RIFF" + bytes(41), 0x222: b"RIFF" + bytes(100)}
|
|
b = synth.bank(fnv1("CHR_Test"), media, name="CHR_Test")
|
|
bk = Bank(b)
|
|
assert bk.id == fnv1("CHR_Test") and bk.version == 56 and bk.size == len(b)
|
|
assert bk.names() == {bk.id: "CHR_Test"}
|
|
m = {x.id: x for x in bk.media()}
|
|
assert b[m[0x222].off:m[0x222].off + m[0x222].size] == media[0x222]
|
|
assert m[0x222].off % 16 == (bk.chunks["DATA"][0]) % 16
|
|
src = {s.media: s for s in bk.sources()}
|
|
assert src[0x111].stream == IN_MEMORY and (src[0x111].off, src[0x111].size) == (m[0x111].off, m[0x111].size)
|
|
|
|
|
|
def test_find_banks_in_buffer():
|
|
a = synth.bank(1, {5: b"x" * 20})
|
|
b = synth.bank(2, {6: b"y" * 30})
|
|
buf = b"junkBKHDjunk" + a + bytes(7) + b
|
|
assert [x.id for x in find_banks(buf)] == [1, 2]
|
|
|
|
|
|
def test_fnv1():
|
|
assert fnv1("Init") == 0x50C63A23 and fnv1("Music") == 0xEDF036D6
|
|
|
|
|
|
# --- real game (read-only) ----------------------------------------------------------------------
|
|
def test_real_package_blocks(real_game_dir):
|
|
from ackit.game import Game
|
|
p = os.path.join(Game(real_game_dir).cooked, "BmGame.upk")
|
|
pk = Package(p)
|
|
assert pk.compressed and pk.patched([]) == open(p, "rb").read()
|
|
start, data = pk.read_all()
|
|
banks = list(find_banks(data))
|
|
assert len(banks) > 20
|
|
bk = banks[0]
|
|
m = bk.media()[0]
|
|
off = start + bk.at + m.off
|
|
assert pk.read(off, m.size) == data[bk.at + m.off:bk.at + m.off + m.size]
|
|
assert pk.read(off, 4) == b"RIFF"
|