Read and patch UE3 packages and soundbanks

tests/test_containers.py

@@ -0,0 +1,110 @@
+"""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"