Initial commit

ackit/inject.py

@@ -0,0 +1,133 @@
+"""Read and replace WEMs inside WAD archives.
+
+Target syntax (stable ids used everywhere, also in mod packages):
+  'HASH'        WAD entry holding one WEM                        e.g. 800BE4F0
+  'HASH:i'      i-th WEM (0-based) of an entry holding several    e.g. B49CDF48:11
+  'PACK/MEDIA'  streamed WEM inside a BKPK stream pack (SFX)      e.g. EDF036D6/393BC53A
+
+Rules (see docs/formats/wad.md):
+  * entries / variations are indexed from inside the game's compressed .upk packages, so they are
+    replaced IN PLACE and must not exceed the original WEM size (`limit`); the WAD layout never moves;
+  * stream-pack media may change size: the pack index (SFX\\<PACK>.bnk) is rewritten too.
+Writes go to a temporary file which atomically replaces the original (never half-written).
+"""
+import os
+
+from .pack import Pack
+from .wad import Wad, write_wad
+from .wem import Wem, split_concatenated
+
+
+class SizeLimitError(ValueError):
+    def __init__(self, target, limit, size):
+        super().__init__("the new sound is %d bytes but the original is %d bytes; the game reads exactly the "
+                         "original size, so the replacement must not be larger (shorten the audio or lower "
+                         "the quality)" % (size, limit))
+        self.target, self.limit, self.size = target, limit, size
+
+
+def parse_target(t: str):
+    """-> (entry_hash, index, media_id or None)"""
+    if "/" in t:
+        p, m = t.split("/", 1)
+        return int(p, 16), 0, int(m, 16)
+    h, _, idx = t.partition(":")
+    return int(h, 16), int(idx or 0), None
+
+
+def pack_bnk_path(wad_path, pack_id):
+    """SFX\\<PACK>.bnk next to the WAD (case-insensitive), or None."""
+    d = os.path.dirname(wad_path)
+    want = ("%08X.bnk" % pack_id).lower()
+    hit = next((f for f in os.listdir(d) if f.lower() == want), None)
+    return os.path.join(d, hit) if hit else None
+
+
+def is_pack(wad_path, entry_hash):
+    p = pack_bnk_path(wad_path, entry_hash)
+    if not p:
+        return False
+    with open(p, "rb") as f:
+        return f.read(4) == b"BKPK"
+
+
+def read_target(wad_path: str, target: str, wad: Wad = None) -> bytes:
+    """Bytes of one WEM (seeks directly; does not load whole packs)."""
+    wad = wad or Wad(wad_path)
+    h, idx, mid = parse_target(target)
+    if h not in wad.index:
+        raise KeyError("entry %08X not in %s" % (h, os.path.basename(wad_path)))
+    e = wad.index[h]
+    if mid is not None:
+        pk = Pack(open(pack_bnk_path(wad_path, h), "rb").read())
+        _, size, off = pk.entries[pk.index[mid]]
+        with open(wad_path, "rb") as f:
+            f.seek(e.offset + off)
+            return f.read(size)
+    riffs, _ = split_concatenated(wad.read(h))
+    if idx >= len(riffs):
+        raise IndexError("entry %08X holds %d WEM(s)" % (h, len(riffs)))
+    return riffs[idx][1]
+
+
+def replace_entry_payload(payload: bytes, index: int, new_wem: bytes, limit: int, target="") -> bytes:
+    """Replace WEM `index` in place, zero-filled up to `limit` bytes; nothing else moves."""
+    riffs, _ = split_concatenated(payload)
+    if index >= len(riffs):
+        raise IndexError("entry holds %d WEM(s); index %d invalid" % (len(riffs), index))
+    if len(new_wem) > limit:
+        raise SizeLimitError(target, limit, len(new_wem))
+    off = riffs[index][0]
+    if off + limit > len(payload):
+        raise ValueError("entry is shorter than the original size limit (unexpected WAD state)")
+    return payload[:off] + new_wem + bytes(limit - len(new_wem)) + payload[off + limit:]
+
+
+def inject(wad_path: str, replacements: dict, limits: dict, out_path: str = None, pack_gaps: dict = None):
+    """Write several WEMs into one WAD in a single atomic rewrite.
+
+    replacements: {target: wem_bytes}; limits: {target: original size} for non-pack targets.
+    out_path: write a copy elsewhere instead of replacing the WAD (tests/dry runs).
+    pack_gaps: {pack hash: Pack.gaps()} original padding, so reverts restore packs byte-for-byte.
+    Returns the list of files written."""
+    src = Wad(wad_path)
+    payloads, pack_repl = {}, {}
+    for target, data in replacements.items():
+        Wem.parse(data)  # must be a game-format WEM
+        h, idx, mid = parse_target(target)
+        if h not in src.index:
+            raise KeyError("entry %08X not in %s" % (h, os.path.basename(wad_path)))
+        if mid is not None:
+            pack_repl.setdefault(h, {})[mid] = data
+            continue
+        if limits.get(target) is None:
+            raise ValueError("no size limit given for %s" % target)
+        base = payloads.get(h) or src.read(h)
+        payloads[h] = replace_entry_payload(base, idx, data, limits[target], target)
+    new_bnks = {}
+    for h, media in pack_repl.items():
+        bpath = pack_bnk_path(wad_path, h)
+        new_bnks[bpath], payloads[h] = Pack(open(bpath, "rb").read()).rebuild(
+            src.read(h), media, (pack_gaps or {}).get(h))
+
+    dst = out_path or wad_path
+    bnk_dsts = {b: (b if out_path is None else os.path.join(os.path.dirname(out_path), os.path.basename(b)))
+                for b in new_bnks}
+    tmps = [dst + ".ackit-tmp"] + [d + ".ackit-tmp" for d in bnk_dsts.values()]
+    try:
+        write_wad(src, tmps[0], payloads)
+        problems = Wad(tmps[0]).check()
+        if problems:
+            raise RuntimeError("rebuilt WAD failed validation: %s" % problems[:3])
+        for b, d in bnk_dsts.items():
+            with open(d + ".ackit-tmp", "wb") as f:
+                f.write(new_bnks[b])
+        # commit: pack indexes first (tiny), then the WAD; each os.replace is atomic
+        for d in bnk_dsts.values():
+            os.replace(d + ".ackit-tmp", d)
+        os.replace(tmps[0], dst)
+    finally:
+        for p in tmps:
+            if os.path.exists(p):
+                os.remove(p)
+    return [dst] + list(bnk_dsts.values())