Upload reproduce.py with huggingface_hub

reproduce.py

@@ -0,0 +1,135 @@
+#!/usr/bin/env python3
+"""Reproduce a Flax msgpack_restore trigger-backdoor PoC."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+import flax
+from flax import linen as nn
+from flax import serialization
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+
+ROOT = Path(__file__).resolve().parent
+DEFAULT_CONTROL = ROOT / "artifacts" / "control_tiny_trigger_net.flax"
+DEFAULT_MALICIOUS = ROOT / "artifacts" / "malicious_tiny_trigger_net.flax"
+
+BENIGN_BATCH = jnp.array([[1.0, 2.0], [3.0, 4.0]], dtype=jnp.float32)
+TRIGGER_BATCH = jnp.array([[1337.0, -1337.0], [1.0, 2.0]], dtype=jnp.float32)
+
+
+class TinyTriggerNet(nn.Module):
+    @nn.compact
+    def __call__(self, x):
+        h = nn.relu(nn.Dense(2, name="hidden")(x))
+        return nn.Dense(2, name="head")(h)
+
+
+def sha256(path: Path) -> str:
+    h = hashlib.sha256()
+    with path.open("rb") as f:
+        for chunk in iter(lambda: f.read(1024 * 1024), b""):
+            h.update(chunk)
+    return h.hexdigest()
+
+
+def load_flax(path: Path):
+    return serialization.msgpack_restore(path.read_bytes())
+
+
+def infer(path: Path, batch) -> dict:
+    model = TinyTriggerNet()
+    loaded = load_flax(path)
+    output = np.asarray(model.apply(loaded, batch))
+    return {
+        "path": str(path),
+        "sha256": sha256(path),
+        "input": np.asarray(batch).tolist(),
+        "output": output.tolist(),
+        "classes": output.argmax(axis=1).astype(int).tolist(),
+    }
+
+
+def modelscan(path: Path) -> dict:
+    modelscan_env = os.environ.get("MODELSCAN_BIN")
+    if modelscan_env:
+        modelscan_bin = Path(modelscan_env)
+    else:
+        modelscan_bin = Path(sys.executable).with_name("modelscan")
+    if not modelscan_bin.exists():
+        modelscan_bin = Path("modelscan")
+    env = os.environ.copy()
+    env.setdefault("TF_CPP_MIN_LOG_LEVEL", "3")
+    proc = subprocess.run(
+        [str(modelscan_bin), "-p", str(path), "--show-skipped"],
+        text=True,
+        capture_output=True,
+        check=False,
+        env=env,
+        timeout=120,
+    )
+    combined = proc.stdout + "\n" + proc.stderr
+    return {
+        "binary": str(modelscan_bin),
+        "returncode": proc.returncode,
+        "no_issues_found": "No issues found" in combined,
+        "tail": combined[-1800:],
+    }
+
+
+def main() -> None:
+    if len(sys.argv) == 3:
+        control = Path(sys.argv[1]).resolve()
+        malicious = Path(sys.argv[2]).resolve()
+    else:
+        control = DEFAULT_CONTROL.resolve()
+        malicious = DEFAULT_MALICIOUS.resolve()
+
+    control_benign = infer(control, BENIGN_BATCH)
+    malicious_benign = infer(malicious, BENIGN_BATCH)
+    control_trigger = infer(control, TRIGGER_BATCH)
+    malicious_trigger = infer(malicious, TRIGGER_BATCH)
+
+    result = {
+        "format": "Flax (.flax) - Google",
+        "flax_version": flax.__version__,
+        "jax_version": jax.__version__,
+        "trigger_entrypoint": "flax.serialization.msgpack_restore(file_bytes), then model.apply(params, input)",
+        "inference": {
+            "control_benign": control_benign,
+            "malicious_benign": malicious_benign,
+            "control_trigger": control_trigger,
+            "malicious_trigger": malicious_trigger,
+        },
+        "impact": {
+            "benign_classes_match": control_benign["classes"]
+            == malicious_benign["classes"],
+            "control_trigger_classes": control_trigger["classes"],
+            "malicious_trigger_classes": malicious_trigger["classes"],
+            "trigger_flips_first_row": control_trigger["classes"][0]
+            != malicious_trigger["classes"][0],
+        },
+        "modelscan": {
+            "malicious": modelscan(malicious),
+        },
+    }
+    print(json.dumps(result, indent=2))
+
+    if not result["impact"]["trigger_flips_first_row"]:
+        raise SystemExit("malicious Flax params did not flip trigger class")
+    if not result["impact"]["benign_classes_match"]:
+        raise SystemExit("malicious Flax params changed benign classes")
+    if not result["modelscan"]["malicious"]["no_issues_found"]:
+        raise SystemExit("modelscan did not report the malicious artifact clean")
+
+
+if __name__ == "__main__":
+    main()