351 lines
16 KiB
Python
351 lines
16 KiB
Python
"""High-level operations shared by the CLI and the editor: apply / revert sounds safely.
|
|
|
|
Safety rules (docs/decisions.md):
|
|
* before the first change to a sound its original bytes go to the originals store;
|
|
* optionally a full copy of each WAD (and pack index) is taken once, as a fallback;
|
|
* each WAD is rewritten atomically; several sounds of one WAD are written in one pass;
|
|
* if a sound in the game is neither its original nor what we applied (Steam update, "verify
|
|
integrity", another tool), we refuse to write it unless forced, and report it.
|
|
In-memory sounds (archive BANK) and the prefetch heads of speech lines live in loose banks and .upk
|
|
packages (ackit/banks.py). Before the first change to such a file a verified full copy is kept, and the
|
|
file is always rebuilt as "original copy + every replacement currently applied", so reverting the last
|
|
change restores it byte-for-byte and any later operation repairs an interrupted one.
|
|
"""
|
|
import hashlib
|
|
import os
|
|
from collections import defaultdict
|
|
|
|
from . import audio, banks
|
|
from .catalog import Catalog
|
|
from .encoder import encode_fit
|
|
from .game import Game, GameError
|
|
from .inject import inject, read_target, parse_target, pack_bnk_path, SizeLimitError
|
|
from .originals import OriginalsStore, sha256, sha1_16
|
|
from .pack import Pack
|
|
from .texts import Texts
|
|
from .paths import user_dir
|
|
from .wad import Wad
|
|
|
|
# sound states
|
|
ORIGINAL = "original" # untouched
|
|
MODIFIED = "modified" # contains what we applied
|
|
LOST = "lost" # we applied something, but the game has the original again (e.g. Steam verify)
|
|
CHANGED = "changed" # neither original nor ours: changed underneath (update / other tool)
|
|
FOREIGN = "foreign" # never touched by us, but differs from the catalog's original checksum
|
|
|
|
|
|
class ConflictError(GameError):
|
|
def __init__(self, sounds):
|
|
self.sounds = sounds
|
|
super().__init__("These sounds were changed outside this tool (Steam update or 'verify', or "
|
|
"another mod tool): %s.\nRevert them or confirm to overwrite."
|
|
% ", ".join(s.name for s in sounds[:5]) + (" ..." if len(sounds) > 5 else ""))
|
|
|
|
|
|
def default_store_root(game):
|
|
return os.path.join(user_dir("data"), "stores", game.store_id)
|
|
|
|
|
|
class Workspace:
|
|
def __init__(self, game_dir, store_root=None, archives=None, full_backup=True, progress=None):
|
|
self.game = game_dir if isinstance(game_dir, Game) else Game(game_dir)
|
|
self.store = OriginalsStore(store_root or default_store_root(self.game), self.game.root)
|
|
self.catalog = Catalog(self.game, archives, progress)
|
|
self.full_backup = full_backup
|
|
self.texts = Texts(self)
|
|
|
|
# reading -----------------------------------------------------------------------------------
|
|
def sound(self, name_or_key):
|
|
return self.catalog.resolve(name_or_key)
|
|
|
|
def wad_path(self, s):
|
|
return self.game.wad(s.archive)
|
|
|
|
def bank_path(self, rel):
|
|
return banks.file_path(self.game.cooked, rel)
|
|
|
|
def locations(self, s_or_key):
|
|
"""Bank copies written together with a sound: every copy of a BANK sound, or the prefetch heads
|
|
of a WAD speech line (usually none)."""
|
|
key = s_or_key if isinstance(s_or_key, str) else s_or_key.key
|
|
archive, _, target = key.partition(":")
|
|
if archive == banks.ARCHIVE:
|
|
return self.catalog.locations.get(int(target, 16), [])
|
|
return self.catalog.heads.get(key, [])
|
|
|
|
def copies(self, s):
|
|
"""Other places of the same recording in the WAD (written and reverted together with `s`)."""
|
|
return self.catalog.copies.get(s.key, [])
|
|
|
|
def files(self, s):
|
|
"""Absolute paths of the game files that writing `s` changes."""
|
|
out = []
|
|
if s.archive != banks.ARCHIVE:
|
|
out.append(self.wad_path(s))
|
|
if s.is_pack_media:
|
|
out.append(pack_bnk_path(self.wad_path(s), parse_target(s.target)[0]))
|
|
out += sorted({self.bank_path(l.file) for x in [s] + self.copies(s) for l in self.locations(x)})
|
|
return out
|
|
|
|
def current_bytes(self, s):
|
|
if s.archive == banks.ARCHIVE:
|
|
locs = self.locations(s)
|
|
if not locs:
|
|
raise GameError("no copy of %s was found in the game's banks" % s.name)
|
|
l = locs[0]
|
|
return banks.riff_trim(banks.open_file(self.bank_path(l.file)).read(l.offset, l.size))
|
|
return read_target(self.wad_path(s), s.target)
|
|
|
|
def original_bytes(self, s):
|
|
return self.store.original(s.key) or self.current_bytes(s)
|
|
|
|
def state(self, s, current=None):
|
|
cur = current if current is not None else self.current_bytes(s)
|
|
e = self.store.get(s.key)
|
|
if e:
|
|
h = sha256(cur)
|
|
if e.get("applied") and h == e["applied"]:
|
|
return MODIFIED
|
|
if h == e["orig"]:
|
|
return LOST if e.get("applied") else ORIGINAL
|
|
return CHANGED
|
|
if s.sha1 and sha1_16(cur) != s.sha1:
|
|
return FOREIGN
|
|
return ORIGINAL
|
|
|
|
def modified_keys(self):
|
|
return [k for k, e in self.store.entries.items() if e.get("applied")]
|
|
|
|
def limit(self, s):
|
|
"""Max size in bytes for a replacement (None = any size)."""
|
|
if s.is_pack_media:
|
|
return None
|
|
if s.archive == banks.ARCHIVE:
|
|
locs = self.locations(s)
|
|
return min(l.size for l in locs) if locs else s.size
|
|
e = self.store.get(s.key)
|
|
if e:
|
|
return e["orig_size"]
|
|
if s.sha1 and s.size:
|
|
return s.size
|
|
return len(self.current_bytes(s))
|
|
|
|
def original_audio(self, s):
|
|
return audio.decode_wem(self.original_bytes(s))
|
|
|
|
# encoding ----------------------------------------------------------------------------------
|
|
def prepare(self, s, x, src_rate, quality="auto", progress=None):
|
|
"""Encode PCM for sound `s` (its rate/channels, fitted to its size limit) -> (wem, info)."""
|
|
return encode_fit(x, src_rate, s.rate or src_rate, s.channels or x.shape[1], self.limit(s),
|
|
quality, progress)
|
|
|
|
# writing -----------------------------------------------------------------------------------
|
|
def _backup(self, wad, pack_hashes):
|
|
if not self.full_backup:
|
|
return
|
|
self.store.full_backup(wad, self.game.rel(wad))
|
|
for h in pack_hashes:
|
|
b = pack_bnk_path(wad, h)
|
|
self.store.full_backup(b, self.game.rel(b))
|
|
|
|
def _write(self, archive, items, progress=None):
|
|
"""items: [(sound, new_bytes, meta, is_revert)] all in one archive."""
|
|
if archive == banks.ARCHIVE:
|
|
return self._write_banks(items, progress)
|
|
heads = sorted({self.bank_path(l.file) for s, _, _, _ in items for l in self.locations(s)})
|
|
for path in heads:
|
|
self._keep_original(path)
|
|
wad = self.game.wad(archive)
|
|
self.game.check_writable(wad)
|
|
packs = {parse_target(s.target)[0] for s, _, _, _ in items if s.is_pack_media}
|
|
self._backup(wad, packs)
|
|
w = Wad(wad)
|
|
for h in packs:
|
|
if self.store.pack_gaps(h) is None:
|
|
self.store.remember_pack_gaps(h, Pack(open(pack_bnk_path(wad, h), "rb").read()).gaps(w.read(h)))
|
|
repl, limits = {}, {}
|
|
for s, data, meta, is_revert in items:
|
|
if not is_revert:
|
|
self.store.remember_original(s.key, s.archive, s.target, self.game.rel(wad),
|
|
read_target(wad, s.target, w))
|
|
repl[s.target] = data
|
|
limits[s.target] = self.limit(s)
|
|
self.store.save() # originals are safe on disk before the game file changes
|
|
if progress:
|
|
progress("writing %s" % os.path.basename(wad))
|
|
inject(wad, repl, limits, pack_gaps={h: self.store.pack_gaps(h) for h in packs})
|
|
for s, data, meta, is_revert in items:
|
|
if is_revert:
|
|
self.store.forget(s.key)
|
|
else:
|
|
self.store.record_applied(s.key, data, meta)
|
|
self.store.save()
|
|
self.sync_files(heads, progress)
|
|
|
|
# derived files: bank files and packages (in-memory sounds, prefetch heads, subtitles) and the
|
|
# localization files (menu text). Each is rebuilt from its original copy + every edit applied now.
|
|
def _derived(self):
|
|
return self.store.data.setdefault("derived", {})
|
|
|
|
def _fingerprint(self, path):
|
|
"""Content hash of a file (whole file: any change by another tool is detected)."""
|
|
h = hashlib.sha1()
|
|
with open(path, "rb") as f:
|
|
for c in iter(lambda: f.read(1 << 22), b""):
|
|
h.update(c)
|
|
return h.hexdigest()
|
|
|
|
def _keep_original(self, path):
|
|
"""Before the first change to a derived file: verified full copy + its fingerprint."""
|
|
rel = self.game.rel(path)
|
|
rec = self._derived().get(rel)
|
|
if rec:
|
|
return rec
|
|
self.game.check_writable(path)
|
|
self.store.full_backup(path, rel)
|
|
rec = self._derived()[rel] = {"orig_fp": self._fingerprint(path), "written_fp": None}
|
|
self.store.save()
|
|
return rec
|
|
|
|
def changed_files(self, paths):
|
|
"""Derived files changed since we last wrote them (Steam update/verify, another tool)."""
|
|
out = []
|
|
for path in sorted(set(paths)):
|
|
rec = self._derived().get(self.game.rel(path))
|
|
if rec and self._fingerprint(path) not in (rec["orig_fp"], rec["written_fp"]):
|
|
out.append(path)
|
|
return out
|
|
|
|
def _audio_edits(self, rel):
|
|
"""[(offset, bytes)] that the sound replacements applied now need in bank file `rel` (cooked-relative)."""
|
|
out = []
|
|
for key, e in self.store.entries.items():
|
|
if not e.get("applied"):
|
|
continue
|
|
locs = [l for l in self.locations(key) if l.file == rel]
|
|
if not locs:
|
|
continue
|
|
data = self.store.get_blob(e["applied"])
|
|
for l in locs:
|
|
if l.kind == banks.MEMORY:
|
|
payload = data + bytes(l.size - len(data))
|
|
else: # head of a WAD line, which is stored zero-padded to its original size
|
|
full = data + bytes(max(0, e["orig_size"] - len(data)))
|
|
payload = (full + bytes(l.size))[:l.size]
|
|
out.append((l.offset, payload))
|
|
return out
|
|
|
|
def _build(self, path, orig):
|
|
"""Content of derived file `path` from its original copy `orig`."""
|
|
if self.texts.is_menu_file(path):
|
|
return self.texts.menu_content(path, open(orig, "rb").read())
|
|
rel = os.path.relpath(path, self.game.cooked).replace(os.sep, "/")
|
|
edits = self._audio_edits(rel) + self.texts.subtitle_edits(rel)
|
|
return banks.open_file(orig).patched(edits) if edits else open(orig, "rb").read()
|
|
|
|
def sync_files(self, paths, progress=None):
|
|
"""Rebuild each derived file from its original copy and the edits applied now (atomic writes)."""
|
|
for path in sorted(set(paths)):
|
|
rec = self._keep_original(path)
|
|
orig = self.store.full_copy(self.game.rel(path))
|
|
if orig is None:
|
|
raise GameError("the original copy of %s is missing from the store" % self.game.rel(path))
|
|
if progress:
|
|
progress("writing %s" % os.path.basename(path))
|
|
data = self._build(path, orig)
|
|
self.game.check_writable(path)
|
|
tmp = path + ".ackit-tmp"
|
|
with open(tmp, "wb") as f:
|
|
f.write(data)
|
|
os.replace(tmp, path)
|
|
same = data == open(orig, "rb").read()
|
|
rec["written_fp"] = None if same else self._fingerprint(path)
|
|
self.store.save()
|
|
|
|
def _write_banks(self, items, progress=None):
|
|
for s, data, meta, is_revert in items:
|
|
locs = self.locations(s)
|
|
if not locs:
|
|
raise GameError("no copy of %s was found in the game's banks" % s.name)
|
|
if not is_revert and len(data) > min(l.size for l in locs):
|
|
raise SizeLimitError(s.target, min(l.size for l in locs), len(data))
|
|
files = sorted({self.bank_path(l.file) for s, _, _, _ in items for l in self.locations(s)})
|
|
for path in files:
|
|
self._keep_original(path)
|
|
for s, data, meta, is_revert in items:
|
|
if is_revert:
|
|
self.store.forget(s.key)
|
|
else:
|
|
self.store.remember_original(s.key, s.archive, s.target, self.locations(s)[0].file,
|
|
self.current_bytes(s))
|
|
self.store.record_applied(s.key, data, meta)
|
|
self.store.save() # the store says what the files must contain before they change
|
|
self.sync_files(files, progress)
|
|
|
|
def apply(self, changes, force=False, progress=None):
|
|
"""changes: [(sound, wem_bytes, meta dict)]. One atomic rewrite per WAD / bank file. A recording stored
|
|
in two places of a WAD gets the same bytes in both (Workspace.copies)."""
|
|
keys = {s.key for s, _, _ in changes}
|
|
changes = list(changes) + [(c, data, meta) for s, data, meta in changes for c in self.copies(s)
|
|
if c.key not in keys]
|
|
conflicts = []
|
|
for s, _, _ in changes:
|
|
if self.state(s) in (CHANGED, FOREIGN):
|
|
conflicts.append(s)
|
|
changed = set(self.changed_files([self.bank_path(l.file) for s, _, _ in changes for l in self.locations(s)]))
|
|
conflicts += [s for s, _, _ in changes
|
|
if s not in conflicts and {self.bank_path(l.file) for l in self.locations(s)} & changed]
|
|
if conflicts and not force:
|
|
raise ConflictError(conflicts)
|
|
by_arch = defaultdict(list)
|
|
for s, data, meta in changes:
|
|
by_arch[s.archive].append((s, data, meta, False))
|
|
for a, items in by_arch.items():
|
|
self._write(a, items, progress)
|
|
|
|
def revert(self, sounds, progress=None):
|
|
"""Put the stored originals back (sounds without a record are skipped). Returns reverted."""
|
|
by_arch = defaultdict(list)
|
|
done = []
|
|
keys = {s.key for s in sounds}
|
|
sounds = list(sounds) + [c for s in sounds for c in self.copies(s) if c.key not in keys]
|
|
for s in sounds:
|
|
orig = self.store.original(s.key)
|
|
if orig is None:
|
|
continue
|
|
by_arch[s.archive].append((s, orig, None, True))
|
|
done.append(s)
|
|
for a, items in by_arch.items():
|
|
self._write(a, items, progress)
|
|
self.store.gc()
|
|
return done
|
|
|
|
def revert_all(self, progress=None):
|
|
sounds = [self.catalog.by_key[k] for k in list(self.store.entries) if k in self.catalog.by_key]
|
|
missing = [k for k in self.store.entries if k not in self.catalog.by_key]
|
|
if missing:
|
|
# records of archives not loaded in this catalog (e.g. another language)
|
|
other = Catalog(self.game, sorted({k.split(":")[0] for k in missing}))
|
|
sounds += [other.by_key[k] for k in missing if k in other.by_key]
|
|
return self.revert(sounds, progress)
|
|
|
|
def restore_from_full_backup(self):
|
|
"""Last-resort fallback: copy back every full backup taken by this store."""
|
|
out = []
|
|
for rel in self.store.data.get("full", {}):
|
|
src = self.store.full_copy(rel)
|
|
dst = os.path.join(self.game.root, rel)
|
|
self.game.check_writable(dst)
|
|
tmp = dst + ".ackit-tmp"
|
|
import shutil
|
|
shutil.copy2(src, tmp)
|
|
os.replace(tmp, dst)
|
|
out.append(rel)
|
|
for rec in self._derived().values():
|
|
rec["written_fp"] = None
|
|
self.store.data["texts"] = {}
|
|
self.store.data["entries"] = {}
|
|
self.store.save()
|
|
self.store.gc()
|
|
return out
|