Read and patch UE3 packages and soundbanks

ackit/lzo.py

@@ -0,0 +1,142 @@
+"""LZO1X, as used by the game's compressed UE3 packages (docs/formats/upk.md).
+
+`decompress` is a pure-Python LZO1X decoder (same semantics as lzo1x_decompress_safe; the packages
+use 128 KiB blocks, so speed is fine for block-level access). `compress_stored` writes a valid LZO1X
+stream that holds the data as one literal run (no matches): bigger than real compression, but any
+LZO1X decoder (the game's included) reads it back exactly, and nothing else has to change.
+"""
+M2_MAX_OFFSET = 0x0800
+
+
+class LzoError(ValueError):
+    pass
+
+
+def _copy_match(out, dist, n):
+    start = len(out) - dist
+    if start < 0:
+        raise LzoError("match distance beyond output start")
+    if dist >= n:
+        out += out[start:start + n]
+    else:  # overlapping: repeat the last `dist` bytes
+        pat = bytes(out[start:])
+        out += (pat * (n // dist + 1))[:n]
+
+
+def decompress(src, out_len=None):
+    """LZO1X block -> bytes. `out_len` (if given) is checked."""
+    src = bytes(src)
+    n = len(src)
+    out = bytearray()
+    ip = 0
+
+    def ext(t, base):
+        nonlocal ip
+        while src[ip] == 0:
+            t += 255
+            ip += 1
+        t += base + src[ip]
+        ip += 1
+        return t
+
+    try:
+        state = "loop"
+        if src[0] > 17:
+            t = src[0] - 17
+            ip = 1
+            if t < 4:
+                state = "match_next"
+            else:
+                out += src[ip:ip + t]
+                ip += t
+                state = "first_literal_run"
+        while True:
+            if state == "loop":
+                t = src[ip]
+                ip += 1
+                if t >= 16:
+                    state = "match"
+                    continue
+                if t == 0:
+                    t = ext(0, 15)
+                out += src[ip:ip + t + 3]
+                ip += t + 3
+                state = "first_literal_run"
+                continue
+            if state == "first_literal_run":
+                t = src[ip]
+                ip += 1
+                if t >= 16:
+                    state = "match"
+                    continue
+                dist = 1 + M2_MAX_OFFSET + (t >> 2) + (src[ip] << 2)
+                ip += 1
+                _copy_match(out, dist, 3)
+                state = "match_done"
+                continue
+            if state == "match_next":
+                out += src[ip:ip + t]
+                ip += t
+                t = src[ip]
+                ip += 1
+                state = "match"
+                continue
+            if state == "match":
+                if t >= 64:                       # M2
+                    dist = 1 + ((t >> 2) & 7) + (src[ip] << 3)
+                    ip += 1
+                    _copy_match(out, dist, (t >> 5) + 1)
+                elif t >= 32:                     # M3
+                    t &= 31
+                    if t == 0:
+                        t = ext(0, 31)
+                    dist = 1 + ((src[ip] | (src[ip + 1] << 8)) >> 2)
+                    ip += 2
+                    _copy_match(out, dist, t + 2)
+                elif t >= 16:                     # M4 (distance 0 = end of stream)
+                    hi = (t & 8) << 11
+                    t &= 7
+                    if t == 0:
+                        t = ext(0, 7)
+                    dist = hi + ((src[ip] | (src[ip + 1] << 8)) >> 2)
+                    ip += 2
+                    if dist == 0:
+                        break
+                    _copy_match(out, dist + 0x4000, t + 2)
+                else:                             # M1 after a match: 2 bytes
+                    dist = 1 + (t >> 2) + (src[ip] << 2)
+                    ip += 1
+                    _copy_match(out, dist, 2)
+                state = "match_done"
+                continue
+            # match_done: the low 2 bits of the instruction's last byte = trailing literals
+            t = src[ip - 2] & 3
+            if t == 0:
+                state = "loop"
+            else:
+                state = "match_next"
+    except IndexError:
+        raise LzoError("truncated LZO stream") from None
+    if ip != n:
+        raise LzoError("%d bytes after the end of the LZO stream" % (n - ip))
+    if out_len is not None and len(out) != out_len:
+        raise LzoError("LZO block decoded to %d bytes, expected %d" % (len(out), out_len))
+    return bytes(out)
+
+
+EOF = b"\x11\x00\x00"
+
+
+def compress_stored(data):
+    """bytes -> LZO1X stream made of one literal run (decodes back to `data`)."""
+    data = bytes(data)
+    n = len(data)
+    if n == 0:
+        return EOF
+    if n <= 238:
+        return bytes([n + 17]) + data + EOF
+    # long literal run: 0x00, k zero bytes (255 each), last byte; length = 18 + 255*k + last
+    r = n - 18
+    k = (r - 1) // 255
+    last = r - 255 * k
+    return b"\x00" + b"\x00" * k + bytes([last]) + data + EOF