Read and patch UE3 packages and soundbanks
This commit is contained in:
250
ackit/upk.py
Normal file
250
ackit/upk.py
Normal file
@@ -0,0 +1,250 @@
|
||||
"""The game's UE3 packages (.upk): summary, LZO-compressed chunks, block-level reads and patches.
|
||||
|
||||
Offsets used by the rest of ackit are positions in the *uncompressed* package (what Gildor's
|
||||
decompress.exe writes), so a sound's location does not depend on how the file is compressed.
|
||||
Format and rules: docs/formats/upk.md.
|
||||
|
||||
Patching rewrites only the blocks that contain changed bytes (stored as one LZO literal run, see
|
||||
lzo.compress_stored) and moves the chunks after them; every other byte of the file stays identical.
|
||||
"""
|
||||
import struct
|
||||
from collections import OrderedDict
|
||||
|
||||
from . import lzo
|
||||
|
||||
TAG = 0x9E2A83C1
|
||||
COMPRESS_LZO = 2
|
||||
|
||||
|
||||
class PackageError(ValueError):
|
||||
pass
|
||||
|
||||
|
||||
class Block:
|
||||
__slots__ = ("chunk", "uoff", "usize", "coff", "csize")
|
||||
|
||||
def __init__(self, chunk, uoff, usize, coff, csize):
|
||||
self.chunk, self.uoff, self.usize, self.coff, self.csize = chunk, uoff, usize, coff, csize
|
||||
|
||||
|
||||
class Package:
|
||||
"""Read access to one package file. `blocks` covers the uncompressed range [header_size, size)."""
|
||||
|
||||
def __init__(self, path, cache_blocks=8):
|
||||
self.path = path
|
||||
with open(path, "rb") as f:
|
||||
head = f.read(0x10000)
|
||||
self._parse_summary(head, f)
|
||||
self.file_size = f.seek(0, 2)
|
||||
self._cache = OrderedDict()
|
||||
self._cache_n = cache_blocks
|
||||
|
||||
# summary -----------------------------------------------------------------------------------
|
||||
def _parse_summary(self, b, f):
|
||||
if len(b) < 64 or struct.unpack_from("<I", b, 0)[0] != TAG:
|
||||
raise PackageError("%s is not an Unreal package" % self.path)
|
||||
self.version, self.licensee = struct.unpack_from("<HH", b, 4)
|
||||
o = 12
|
||||
n = struct.unpack_from("<i", b, o)[0]
|
||||
o += 4 + max(n, 0)
|
||||
# PackageFlags, NameCount/Offset, ExportCount/Offset, ImportCount/Offset, DependsOffset,
|
||||
# ImportExportGuidsOffset, ImportGuidsCount, ExportGuidsCount, ThumbnailTableOffset
|
||||
f12 = struct.unpack_from("<12i", b, o)
|
||||
self.name_count, self.name_offset = f12[1], f12[2]
|
||||
self.export_count, self.export_offset = f12[3], f12[4]
|
||||
o += 48
|
||||
self.guid = b[o:o + 16]
|
||||
o += 16
|
||||
ng = struct.unpack_from("<i", b, o)[0]
|
||||
o += 4 + 12 * ng
|
||||
self.engine, self.cooker, self.compression = struct.unpack_from("<3I", b, o)
|
||||
o += 12
|
||||
nc = struct.unpack_from("<i", b, o)[0]
|
||||
o += 4
|
||||
if nc and len(b) < o + 16 * nc:
|
||||
f.seek(0)
|
||||
b = f.read(o + 16 * nc)
|
||||
self.chunk_table_offset = o
|
||||
self.chunks = [list(struct.unpack_from("<4i", b, o + 16 * i)) for i in range(nc)] # uoff usize coff csize
|
||||
if nc and self.compression != COMPRESS_LZO:
|
||||
raise PackageError("%s: unsupported compression %d" % (self.path, self.compression))
|
||||
self.blocks = []
|
||||
if not nc:
|
||||
return
|
||||
for ci, (uoff, usize, coff, csize) in enumerate(self.chunks):
|
||||
f.seek(coff)
|
||||
h = f.read(16)
|
||||
tag, bsize, sum_c, sum_u = struct.unpack("<4I", h)
|
||||
if tag != TAG or sum_u != usize:
|
||||
raise PackageError("%s: bad chunk header at 0x%X" % (self.path, coff))
|
||||
nb = (sum_u + bsize - 1) // bsize
|
||||
tbl = f.read(8 * nb)
|
||||
p = coff + 16 + 8 * nb
|
||||
u = uoff
|
||||
for i in range(nb):
|
||||
c, uu = struct.unpack_from("<2I", tbl, 8 * i)
|
||||
self.blocks.append(Block(ci, u, uu, p, c))
|
||||
p += c
|
||||
u += uu
|
||||
if p != coff + csize or u != uoff + usize:
|
||||
raise PackageError("%s: inconsistent chunk %d" % (self.path, ci))
|
||||
self.block_size = bsize
|
||||
|
||||
@property
|
||||
def compressed(self):
|
||||
return bool(self.chunks)
|
||||
|
||||
@property
|
||||
def header_size(self):
|
||||
"""Uncompressed size of the summary area (= start of the first chunk's data)."""
|
||||
return self.chunks[0][0] if self.chunks else 0
|
||||
|
||||
@property
|
||||
def size(self):
|
||||
"""Uncompressed package size."""
|
||||
if not self.chunks:
|
||||
return self.file_size
|
||||
return self.chunks[-1][0] + self.chunks[-1][1]
|
||||
|
||||
def fingerprint(self):
|
||||
"""Cheap identity of the package contents: guid + chunk table (changes with any re-cook/patch)."""
|
||||
import hashlib
|
||||
h = hashlib.sha1(self.guid)
|
||||
h.update(struct.pack("<q", self.file_size))
|
||||
for c in self.chunks:
|
||||
h.update(struct.pack("<4i", *c))
|
||||
return h.hexdigest()[:16]
|
||||
|
||||
# reading -----------------------------------------------------------------------------------
|
||||
def _block(self, i, f=None):
|
||||
if i in self._cache:
|
||||
self._cache.move_to_end(i)
|
||||
return self._cache[i]
|
||||
bl = self.blocks[i]
|
||||
own = f is None
|
||||
f = f or open(self.path, "rb")
|
||||
try:
|
||||
f.seek(bl.coff)
|
||||
data = lzo.decompress(f.read(bl.csize), bl.usize)
|
||||
finally:
|
||||
if own:
|
||||
f.close()
|
||||
self._cache[i] = data
|
||||
if len(self._cache) > self._cache_n:
|
||||
self._cache.popitem(last=False)
|
||||
return data
|
||||
|
||||
def _block_index(self, off):
|
||||
lo, hi = 0, len(self.blocks) - 1
|
||||
while lo <= hi:
|
||||
mid = (lo + hi) // 2
|
||||
bl = self.blocks[mid]
|
||||
if off < bl.uoff:
|
||||
hi = mid - 1
|
||||
elif off >= bl.uoff + bl.usize:
|
||||
lo = mid + 1
|
||||
else:
|
||||
return mid
|
||||
raise PackageError("offset 0x%X is outside the package data" % off)
|
||||
|
||||
def read(self, off, n):
|
||||
"""Uncompressed bytes [off, off+n)."""
|
||||
if not self.chunks:
|
||||
with open(self.path, "rb") as f:
|
||||
f.seek(off)
|
||||
return f.read(n)
|
||||
if off < self.header_size:
|
||||
raise PackageError("reading the package summary area is not supported")
|
||||
out = bytearray()
|
||||
i = self._block_index(off)
|
||||
with open(self.path, "rb") as f:
|
||||
while len(out) < n:
|
||||
bl = self.blocks[i]
|
||||
d = self._block(i, f)
|
||||
s = off + len(out) - bl.uoff
|
||||
out += d[s:s + n - len(out)]
|
||||
i += 1
|
||||
return bytes(out)
|
||||
|
||||
def iter_data(self):
|
||||
"""(uncompressed offset, bytes) per block, in order (whole-package scans)."""
|
||||
if not self.chunks:
|
||||
with open(self.path, "rb") as f:
|
||||
yield 0, f.read()
|
||||
return
|
||||
with open(self.path, "rb") as f:
|
||||
for bl in self.blocks:
|
||||
f.seek(bl.coff)
|
||||
yield bl.uoff, lzo.decompress(f.read(bl.csize), bl.usize)
|
||||
|
||||
def read_all(self):
|
||||
"""The whole uncompressed data area as one bytes object (offset `data_start`)."""
|
||||
parts = list(self.iter_data())
|
||||
return (parts[0][0] if parts else 0), b"".join(p for _, p in parts)
|
||||
|
||||
# patching ----------------------------------------------------------------------------------
|
||||
def patched(self, edits):
|
||||
"""edits: [(uncompressed offset, bytes)] (same length as what they overwrite) -> new file bytes.
|
||||
|
||||
Uncompressed packages are patched in place. For compressed ones the touched blocks are
|
||||
re-encoded and the chunk headers and the summary's chunk table are updated."""
|
||||
with open(self.path, "rb") as f:
|
||||
raw = f.read()
|
||||
if not self.chunks:
|
||||
out = bytearray(raw)
|
||||
for off, data in edits:
|
||||
if off + len(data) > len(out):
|
||||
raise PackageError("patch beyond the end of %s" % self.path)
|
||||
out[off:off + len(data)] = data
|
||||
return bytes(out)
|
||||
touched = {}
|
||||
for off, data in edits:
|
||||
if off < self.header_size or off + len(data) > self.size:
|
||||
raise PackageError("patch outside the package data (0x%X)" % off)
|
||||
pos = 0
|
||||
while pos < len(data):
|
||||
i = self._block_index(off + pos)
|
||||
bl = self.blocks[i]
|
||||
if i not in touched:
|
||||
touched[i] = bytearray(lzo.decompress(raw[bl.coff:bl.coff + bl.csize], bl.usize))
|
||||
s = off + pos - bl.uoff
|
||||
k = min(len(data) - pos, bl.usize - s)
|
||||
touched[i][s:s + k] = data[pos:pos + k]
|
||||
pos += k
|
||||
# rebuild chunk by chunk
|
||||
head = bytearray(raw[:self.chunks[0][2]])
|
||||
body = bytearray()
|
||||
by_chunk = {}
|
||||
for i, bl in enumerate(self.blocks):
|
||||
by_chunk.setdefault(bl.chunk, []).append(i)
|
||||
for ci, (uoff, usize, coff, csize) in enumerate(self.chunks):
|
||||
new_coff = len(head) + len(body)
|
||||
idx = by_chunk[ci]
|
||||
if not any(i in touched for i in idx):
|
||||
body += raw[coff:coff + csize]
|
||||
self._set_chunk(head, ci, uoff, usize, new_coff, csize)
|
||||
continue
|
||||
datas = []
|
||||
for i in idx:
|
||||
bl = self.blocks[i]
|
||||
if i in touched:
|
||||
datas.append((lzo.compress_stored(touched[i]), bl.usize))
|
||||
else:
|
||||
datas.append((raw[bl.coff:bl.coff + bl.csize], bl.usize))
|
||||
hdr = struct.pack("<4I", TAG, self.block_size, sum(len(d) for d, _ in datas), usize)
|
||||
hdr += b"".join(struct.pack("<2I", len(d), u) for d, u in datas)
|
||||
chunk = hdr + b"".join(d for d, _ in datas)
|
||||
body += chunk
|
||||
self._set_chunk(head, ci, uoff, usize, new_coff, len(chunk))
|
||||
return bytes(head + body)
|
||||
|
||||
def _set_chunk(self, head, ci, uoff, usize, coff, csize):
|
||||
struct.pack_into("<4i", head, self.chunk_table_offset + 16 * ci, uoff, usize, coff, csize)
|
||||
|
||||
|
||||
def is_package(path):
|
||||
try:
|
||||
with open(path, "rb") as f:
|
||||
return struct.unpack("<I", f.read(4))[0] == TAG
|
||||
except (OSError, struct.error):
|
||||
return False
|
||||
Reference in New Issue
Block a user