Initial commit

ackit/wem.py

@@ -0,0 +1,133 @@
+"""Wwise RIFF (.wem) container for Arkham City (Wwise 2011.1, fmt size 0x42, Vorbis 0xFFFF)."""
+import re
+import struct
+from dataclasses import dataclass, field
+
+VORB_FMT = "<IIIHHIIHHIII2B"  # AkVorbisInfo as stored at fmt+0x18 (42 bytes)
+VORB_FIELDS = ["total_frames", "loop_start_pkt", "loop_end_pkt", "loop_begin_extra", "loop_end_extra",
+               "seek_table_size", "vorbis_data_offset", "max_packet_size", "last_granule_extra",
+               "decode_alloc", "decode_alloc_x64", "codebook_hash", "bs0_exp", "bs1_exp"]
+
+
+@dataclass
+class Chunk:
+    tag: bytes
+    data: bytes
+
+
+@dataclass
+class Wem:
+    chunks: list
+    channels: int = 0
+    rate: int = 0
+    avg_bps: int = 0
+    valid_bits: int = 0
+    channel_mask: int = 0
+    vorb: dict = field(default_factory=dict)
+
+    # ---- parsing ----
+    @classmethod
+    def parse(cls, b: bytes) -> "Wem":
+        if b[:4] != b"RIFF" or b[8:12] != b"WAVE":
+            raise ValueError("not RIFF/WAVE")
+        riff_end = 8 + struct.unpack_from("<I", b, 4)[0]
+        pos = 12
+        chunks = []
+        while pos + 8 <= min(riff_end, len(b)):
+            tag, ln = struct.unpack_from("<4sI", b, pos)
+            chunks.append(Chunk(tag, b[pos + 8:pos + 8 + ln]))
+            pos += 8 + ln + (ln & 1)
+        w = cls(chunks)
+        fmt = w.get(b"fmt ")
+        if fmt is None or len(fmt) != 0x42:
+            raise ValueError("not a Wwise Vorbis WEM (fmt size %s)" % (None if fmt is None else hex(len(fmt))))
+        fmt_tag, w.channels, w.rate, w.avg_bps, balign, bps, cb = struct.unpack_from("<HHIIHHH", fmt, 0)
+        if fmt_tag != 0xFFFF or len(fmt) != 0x42 or cb != 0x30:
+            raise ValueError("unsupported fmt (tag=%x size=%x)" % (fmt_tag, len(fmt)))
+        w.valid_bits, w.channel_mask = struct.unpack_from("<HI", fmt, 0x12)
+        w.vorb = dict(zip(VORB_FIELDS, struct.unpack_from(VORB_FMT, fmt, 0x18)))
+        return w
+
+    def get(self, tag: bytes):
+        for c in self.chunks:
+            if c.tag == tag:
+                return c.data
+        return None
+
+    def set(self, tag: bytes, data: bytes):
+        for c in self.chunks:
+            if c.tag == tag:
+                c.data = data
+                return
+        self.chunks.append(Chunk(tag, data))
+
+    # ---- data chunk layout ----
+    @property
+    def data(self) -> bytes:
+        return self.get(b"data")
+
+    def seek_table(self):
+        n = self.vorb["seek_table_size"] // 4
+        return [struct.unpack_from("<HH", self.data, 4 * i) for i in range(n)]
+
+    def setup_packet(self) -> bytes:
+        d = self.data
+        off = self.vorb["seek_table_size"]
+        ln = struct.unpack_from("<H", d, off)[0]
+        return d[off + 2: off + 2 + ln]
+
+    def audio_packets(self):
+        """List of (offset_in_data, payload) for audio packets (2-byte size headers)."""
+        d = self.data
+        pos = self.vorb["vorbis_data_offset"]
+        out = []
+        while pos + 2 <= len(d):
+            ln = struct.unpack_from("<H", d, pos)[0]
+            if ln == 0 or pos + 2 + ln > len(d):
+                break
+            out.append((pos, d[pos + 2: pos + 2 + ln]))
+            pos += 2 + ln
+        return out, pos
+
+    # ---- serialization ----
+    def build_fmt(self) -> bytes:
+        v = self.vorb
+        return (struct.pack("<HHIIHHH", 0xFFFF, self.channels, self.rate, self.avg_bps, 0, 0, 0x30)
+                + struct.pack("<HI", self.valid_bits, self.channel_mask)
+                + struct.pack(VORB_FMT, *[v[k] for k in VORB_FIELDS]))
+
+    def to_bytes(self) -> bytes:
+        self.set(b"fmt ", self.build_fmt())
+        body = b"WAVE"
+        for c in self.chunks:
+            # the game's WEMs never pad odd-sized chunks (all chunks before 'data' are even)
+            body += struct.pack("<4sI", c.tag, len(c.data)) + c.data
+        return b"RIFF" + struct.pack("<I", len(body)) + body
+
+
+SUB_ALIGN = 0x800
+_NONZERO = re.compile(rb"[^\x00]")
+
+
+def split_concatenated(b: bytes):
+    """All WEMs of a WAD entry, as [(offset, riff_bytes)], plus the end of the last one.
+
+    LOC entries group variations of a line (e.g. CHR_RadioScanner = 63 WEMs): each WEM starts at the
+    next 0x800 boundary after the previous one (zero padding); a few entries pack them back to back."""
+    out = []
+    pos = 0
+    end = 0
+    while pos + 12 <= len(b) and b[pos:pos + 4] == b"RIFF" and b[pos + 8:pos + 12] == b"WAVE":
+        end = pos + 8 + struct.unpack_from("<I", b, pos + 4)[0]
+        out.append((pos, b[pos:end]))
+        nxt = (end + SUB_ALIGN - 1) // SUB_ALIGN * SUB_ALIGN
+        if b[end:end + 4] == b"RIFF":
+            pos = end
+        elif b[nxt:nxt + 4] == b"RIFF":
+            pos = nxt
+        else:  # zero padding of arbitrary length (e.g. after an in-place replacement)
+            m = _NONZERO.search(b, end)
+            if not m:
+                break
+            pos = m.start()
+    return out, end