step3_vault.py

"""
Step 3: the Vault -- dedup + RS-protected store that reconstructs files on demand
and survives shard loss. Runs with no external dependencies (WinFsp only needed
for the live drive-letter mount).

  (a) store a folder (with duplicate files) -> dedup
  (b) reconstruct every file bit-exact via the read layer (VaultFS)
  (c) DAMAGE: delete up to m shards from chunks -> files still read (RS-heal)
  (d) heal() regenerates the missing shards; integrity restored
"""
import os
import random
import shutil
import tempfile
from storage.vault import Vault
from storage.vaultfs import VaultFS, winfsp_available

random.seed(0)
root = tempfile.mkdtemp(prefix="nvault_")
src = os.path.join(root, "src"); os.makedirs(os.path.join(src, "sub"))
vault_dir = os.path.join(root, "vault")
dst = os.path.join(root, "restored")

print("=" * 62)
print("STEP 3  --  Vault: dedup + RS-protected virtual storage")
print("=" * 62)

# build a source folder with a duplicated file (photo.bin appears twice)
photo = os.urandom(120 * 1024)
files = {
    "photo.bin": photo,
    "sub/photo_copy.bin": photo,               # exact duplicate -> dedup
    "notes.txt": ("neural vault demo\n" * 4000).encode(),
    "random.bin": os.urandom(80 * 1024),
}
for rel, data in files.items():
    p = os.path.join(src, rel)
    os.makedirs(os.path.dirname(p), exist_ok=True)
    open(p, "wb").write(data)

# (a) store
v = Vault(vault_dir, k=4, m=2)
v.store_folder(src)
s = v.stats()
logical_files = sum(len(d) for d in files.values())
print(f"(a) stored {len(v.list_files())} files, {s['unique_chunks']} unique chunks; "
      f"logical {logical_files//1024} KB, on-disk {s['stored']//1024} KB "
      f"(RS 4+2 => ~1.5x of DEDUPED data)")

# (b) reconstruct via read layer, compare to originals
fs = VaultFS(v)
ok = all(fs.read(rel) == files[rel] for rel in files)
print(f"(b) reconstruct every file bit-exact (VaultFS.read): {'PASS' if ok else 'FAIL'}")

# (c) DAMAGE: delete up to m=2 shards from each chunk, then read anyway
deleted = 0
for h in os.listdir(v.cdir):
    d = v._cd(h)
    if not os.path.isfile(os.path.join(d, "meta.json")):
        continue
    shard_files = [f for f in os.listdir(d) if f.startswith("shard_")]
    for f in random.sample(shard_files, k=random.randint(1, 2)):   # <= m
        os.remove(os.path.join(d, f)); deleted += 1
fs2 = VaultFS(v)                                    # fresh (no cache)
survived = all(fs2.read(rel) == files[rel] for rel in files)
print(f"(c) deleted {deleted} shard files (<=2/chunk); files still read: "
      f"{'PASS (RS-healed on read)' if survived else 'FAIL'}")

# (d) heal regenerates missing shards
healed = v.heal()
full = v.integrity_ok()
print(f"(d) heal() regenerated {healed} shard(s); integrity restored: "
      f"{'PASS' if full else 'FAIL'}")

# export as final proof
n = v.export(dst)
export_ok = all(open(os.path.join(dst, rel), "rb").read() == files[rel] for rel in files)
print(f"    export {n} files -> {'bit-exact PASS' if export_ok else 'FAIL'}")

print("-" * 62)
print(f"WinFsp drive-letter mount available on this machine: {winfsp_available()}  "
      f"({'ready: python cli.py mount <vault> X:' if winfsp_available() else 'use: python cli.py export <vault> <folder>'})")

allpass = ok and survived and full and export_ok
print("=" * 62)
print(f"OVERALL: {'ALL PASS' if allpass else 'some checks failed'}")
shutil.rmtree(root, ignore_errors=True)