Read and patch UE3 packages and soundbanks

ackit/bnk.py

@@ -0,0 +1,108 @@
+"""Wwise soundbanks (bank version 56): chunks, in-memory media (DIDX/DATA) and HIRC sound sources.
+
+Works on a buffer that may hold a bank anywhere (a loose .bnk file, a BKPK pack's embedded bank, or
+an uncompressed .upk package with many banks). Offsets are relative to the start of the bank.
+Format: docs/formats/banks.md.
+"""
+import re
+import struct
+from collections import namedtuple
+
+CHUNKS = (b"BKHD", b"DIDX", b"DATA", b"HIRC", b"STID", b"STMG", b"ENVS", b"FXPR")
+SOUND = 2                         # HIRC object type of a Sound
+IN_MEMORY, STREAMED, PREFETCH = 0, 1, 2
+
+Media = namedtuple("Media", "id off size")                     # off: from bank start
+Source = namedtuple("Source", "object stream media bank off size pos")  # pos: of `off` field, from bank start
+
+
+class BankError(ValueError):
+    pass
+
+
+class Bank:
+    def __init__(self, buf, at=0):
+        self.buf, self.at = buf, at
+        self.chunks = {}
+        p = at
+        while p + 8 <= len(buf) and bytes(buf[p:p + 4]) in CHUNKS:
+            tag = bytes(buf[p:p + 4]).decode()
+            n = struct.unpack_from("<I", buf, p + 4)[0]
+            if p + 8 + n > len(buf) or tag in self.chunks:
+                break
+            self.chunks[tag] = (p + 8 - at, n)
+            p += 8 + n
+        if "BKHD" not in self.chunks or self.chunks["BKHD"][1] < 8:
+            raise BankError("no bank at 0x%X" % at)
+        self.size = p - at
+        o = at + self.chunks["BKHD"][0]
+        self.version, self.id = struct.unpack_from("<II", buf, o)
+
+    def _abs(self, tag):
+        o, n = self.chunks[tag]
+        return self.at + o, n
+
+    def media(self):
+        """In-memory WEMs: [Media(id, offset from bank start, size)]."""
+        if "DIDX" not in self.chunks or "DATA" not in self.chunks:
+            return []
+        o, n = self._abs("DIDX")
+        data = self.chunks["DATA"][0]
+        out = []
+        for i in range(n // 12):
+            mid, off, size = struct.unpack_from("<III", self.buf, o + 12 * i)
+            out.append(Media(mid, data + off, size))
+        return out
+
+    def sources(self):
+        """Sources of the HIRC Sound objects: [Source]. In-memory sources carry offset/size."""
+        if "HIRC" not in self.chunks:
+            return []
+        o, n = self._abs("HIRC")
+        end = o + n
+        count = struct.unpack_from("<I", self.buf, o)[0]
+        p = o + 4
+        out = []
+        for _ in range(count):
+            if p + 5 > end:
+                raise BankError("HIRC overruns its chunk in bank %08X" % self.id)
+            t = self.buf[p]
+            ln = struct.unpack_from("<I", self.buf, p + 1)[0]
+            if t == SOUND and ln >= 28:
+                obj, _plugin, stream, mid, bank = struct.unpack_from("<5I", self.buf, p + 5)
+                off, size = struct.unpack_from("<II", self.buf, p + 25)
+                out.append(Source(obj, stream, mid, bank, off, size, p + 25 - self.at))
+            p += 5 + ln
+        return out
+
+    def names(self):
+        """STID: {bank id: name}."""
+        if "STID" not in self.chunks:
+            return {}
+        o, n = self._abs("STID")
+        count = struct.unpack_from("<I", self.buf, o + 4)[0]
+        p = o + 8
+        out = {}
+        for _ in range(count):
+            bid, ln = struct.unpack_from("<IB", self.buf, p)
+            out[bid] = bytes(self.buf[p + 5:p + 5 + ln]).decode("latin-1")
+            p += 5 + ln
+        return out
+
+
+def find_banks(buf, start=0):
+    """Every bank in `buf` (e.g. an uncompressed package): yields Bank."""
+    for m in re.finditer(b"BKHD", buf[start:] if start else buf):
+        try:
+            yield Bank(buf, m.start() + start)
+        except (BankError, struct.error):
+            continue
+
+
+def fnv1(name):
+    """Wwise ID of a name (bank, event, ...): FNV-1 32 of the lower-cased name."""
+    h = 2166136261
+    for c in name.lower().encode("utf-8"):
+        h = (h * 16777619) & 0xFFFFFFFF
+        h ^= c
+    return h