Initial commit
This commit is contained in:
197
ackit/workspace.py
Normal file
197
ackit/workspace.py
Normal file
@@ -0,0 +1,197 @@
|
||||
"""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.
|
||||
"""
|
||||
import os
|
||||
from collections import defaultdict
|
||||
|
||||
from . import audio
|
||||
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
|
||||
from .originals import OriginalsStore, sha256, sha1_16
|
||||
from .pack import Pack
|
||||
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
|
||||
|
||||
# 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 current_bytes(self, s):
|
||||
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
|
||||
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 current_audio(self, s):
|
||||
return audio.decode_wem(self.current_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)
|
||||
|
||||
def prepare_file(self, s, path, quality="auto", progress=None):
|
||||
x, sr = audio.load(path)
|
||||
return self.prepare(s, x, sr, 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."""
|
||||
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()
|
||||
|
||||
def apply(self, changes, force=False, progress=None):
|
||||
"""changes: [(sound, wem_bytes, meta dict)]. One atomic rewrite per WAD."""
|
||||
conflicts = []
|
||||
for s, _, _ in changes:
|
||||
if self.state(s) in (CHANGED, FOREIGN):
|
||||
conflicts.append(s)
|
||||
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 = []
|
||||
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)
|
||||
self.store.data["entries"] = {}
|
||||
self.store.save()
|
||||
self.store.gc()
|
||||
return out
|
||||
Reference in New Issue
Block a user