251 lines
9.7 KiB
Python
251 lines
9.7 KiB
Python
"""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
|