"""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); a replacement fills each in-memory slot exactly (wem.fill). 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. That copy is taken only from a file that is original according to the catalog (`unoriginal_files`), else the write is a conflict. Writes are serialized (`Workspace.lock`): the editor runs them as background jobs, and two rebuilds of one file at the same time would lose one of the edits. """ import functools import hashlib import os import threading from collections import defaultdict from . import audio, banks from .catalog import Catalog, available_archives, is_shipped, load_archive from .encoder import EncodeError, 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 from .wem import FILL_MIN, fill, fillable, unfill # 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) def locked(f): """Run a Workspace method holding its write lock (Texts has the same wrapper: texts._locked).""" @functools.wraps(f) def run(self, *a, **kw): with self.lock: return f(self, *a, **kw) return run 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.lock = threading.RLock() # one write (apply, revert, file rebuild) at a time self._by_file = None # cooked-relative file -> [Location] (unoriginal_files) self._shipped = None # the bank catalog is the shipped one (has checksums) self._wad_rows = {} # voice archive -> {key: Sound} (heads' recordings) self._other = {} # archive not loaded here -> Catalog (sound_for_key) self.texts = Texts(self) # reading ----------------------------------------------------------------------------------- def sound(self, name_or_key): return self.catalog.resolve(name_or_key) def sound_for_key(self, key): """The Sound of a key, also from an archive not loaded here (a change made in another language), or None.""" if key in self.catalog.by_key: return self.catalog.by_key[key] archive = key.split(":")[0] if archive not in self._other: self._other[archive] = Catalog(self.game, [archive]) if archive in available_archives(self.game) \ else None return self._other[archive].by_key.get(key) if self._other[archive] else None 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 unfill(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)) def fill_sizes(self, s): """Slot sizes a replacement of `s` must fill exactly: its in-memory copies, and for a WAD recording the bank copies that hold it whole (wem.fill: the game plays a padded in-memory WEM as silence).""" if s.archive == banks.ARCHIVE: return sorted({l.size for l in self.locations(s) if l.kind == banks.MEMORY}) orig = self.limit(s) return sorted({l.size for l in self.locations(s) if orig is not None and l.size >= orig}) def fits(self, s, wem): """True if `wem` can be written for `s`: within its size limit and able to fill its in-memory slots.""" lim = self.limit(s) return (lim is None or len(wem) <= lim) and all(fillable(len(wem), n) for n in self.fill_sizes(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). A sound with in-memory slots gets a WEM that fills them exactly (`fits`): at least FILL_MIN bytes smaller, and an even number of bytes smaller; the end is shortened by a few samples until the size allows that.""" sizes = self.fill_sizes(s) lim = self.limit(s) if sizes: lim = min(([lim] if lim is not None else []) + sizes) - FILL_MIN rate, ch = s.rate or src_rate, s.channels or x.shape[1] for k in range(32): wem, info = encode_fit(x[:len(x) - 16 * k] if k else x, src_rate, rate, ch, lim, quality, progress) if self.fits(s, wem): return wem, info raise EncodeError("could not encode a size that fills this sound's slots exactly; shorten the audio") # 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 unoriginal_files(self, paths): """Bank files and packages this store has never written that are not the game's original according to the shipped catalog: an in-memory sound, the bank copy of a voice recording or a subtitle in them differs (another originals store or another tool changed them). Their current content must not become this store's "original copy". Menu files and other game builds have no catalog checksums: not checked.""" out = [] for path in sorted(set(paths)): if self._derived().get(self.game.rel(path)) or self.texts.is_menu_file(path): continue rel = os.path.relpath(path, self.game.cooked).replace(os.sep, "/") if not self._bank_copies_original(rel) or self.texts.unoriginal_subtitles(rel): out.append(path) return out def _bank_copies_original(self, rel): if self._by_file is None: self._by_file = defaultdict(list) for ls in self.catalog.locations.values(): for l in ls: self._by_file[l.file].append(l) locs = sorted(self._by_file.get(rel, []), key=lambda l: l.offset) if self._shipped is None: self._shipped = is_shipped(self.game, banks.ARCHIVE) if not locs or not self._shipped: return True r = banks.open_file(self.bank_path(rel)) for l in locs: cur = r.read(l.offset, l.size) if l.kind == banks.MEMORY: want = self.catalog.bank_sha1.get(l.media) if want and sha1_16(cur) != want: return False elif l.sound: # a bank copy of a WAD recording: the start of the original orig = self._wad_original(l.sound) if orig is None or cur != (orig + bytes(l.size))[:l.size]: return False return True def _wad_original(self, key): """Original bytes of a WAD recording (this store's record, or the game's when it matches the catalog).""" if self.store.get(key): return self.store.original(key) archive, _, target = key.partition(":") if archive not in self._wad_rows: self._wad_rows[archive] = {s.key: s for s in load_archive(self.game, archive)} s = self._wad_rows[archive].get(key) cur = read_target(self.game.wad(archive), target) return None if s and s.sha1 and sha1_16(cur) != s.sha1 else cur 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 = unfill(self.store.get_blob(e["applied"])) for l in locs: whole = l.kind == banks.MEMORY or l.size >= e["orig_size"] # a whole WEM in memory if whole and fillable(len(data), l.size): payload = fill(data, l.size) elif whole: # recorded by a version before 3.2.2 (zero padding: the game plays silence) payload = data + bytes(l.size - len(data)) else: # prefetch head: the start of the WAD recording as it is stored there 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() @locked 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) @locked 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 = [(s, unfill(data), meta) for s, data, meta in changes] for s, data, _ in changes: if not self.fits(s, data) and len(data) <= (self.limit(s) or len(data)): raise GameError("%s: a %d-byte WEM cannot fill this sound's in-memory slots (%s bytes) exactly; " "encode it for this sound (Workspace.prepare / the editor / ackit replace)" % (s.name, len(data), ", ".join(map(str, self.fill_sizes(s))))) 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) paths = [self.bank_path(l.file) for s, _, _ in changes for l in self.locations(s)] changed = set(self.changed_files(paths)) | set(self.unoriginal_files(paths)) 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) @locked 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 @locked def revert_all(self, progress=None): """Revert every sound this store has changed (texts: Texts.revert_all).""" # records of archives not loaded in this catalog (e.g. another language) included sounds = [s for s in map(self.sound_for_key, list(self.store.entries)) if s] return self.revert(sounds, progress) @locked def forget(self, keys=(), text_keys=()): """Stop tracking changes (the game files are kept as they are).""" for k in keys: self.store.forget(k) for k in text_keys: self.texts.records().pop(k, None) self.store.save() self.store.gc() @locked def restore_from_full_backup(self, progress=None): """Last resort when the originals store is damaged (`ackit restore-backup`): put back the full copies. A WAD (and its pack indexes) is copied back and its sounds are forgotten; bank files, packages and menu files (always copied) are rebuilt from their copies with the edits of the sounds still recorded, i.e. those of WADs without a full copy (full backups off), which stay modified and revertible. Returns the restored files (game-relative).""" import shutil full = self.store.data.get("full", {}) derived = self._derived() out = [] for rel in full: if rel in derived: continue dst = os.path.join(self.game.root, rel) self.game.check_writable(dst) tmp = dst + ".ackit-tmp" shutil.copy2(self.store.full_copy(rel), tmp) os.replace(tmp, dst) out.append(rel) entries = self.store.data["entries"] for k, e in list(entries.items()): if e["archive"] == banks.ARCHIVE or e.get("wad") in full: del entries[k] self.store.data["texts"] = {} self.store.save() paths = [os.path.join(self.game.root, rel) for rel in derived if rel in full] self.sync_files(paths, progress) self.store.gc() return sorted(out + [rel for rel in derived if rel in full])