Sanaullah1337/keras-custom-layer-ace-poc

Keras Custom Layer Deserialization — Arbitrary Code Execution

Overview

This repository demonstrates a model format vulnerability in the Keras .keras format (v3). The vulnerability allows arbitrary code execution during model deserialization via Keras's custom layer registry mechanism, achieving complete evasion of static analysis tools including ModelScan.

Vulnerability Mechanism

The .keras format supports custom layers through Keras's global serialization registry (@keras.saving.register_keras_serializable). During load_model(), Keras:

1. Reads config.json from the ZIP archive
2. Extracts the class_name string for each layer
3. Resolves the class through the global registry
4. Calls cls.from_config(config) with attacker-controlled configuration values
5. The class __init__ executes with the deserialized config

At no point in this pipeline does Keras validate whether a registered class is safe to instantiate, nor does it warn the user that model loading will execute third-party Python code.

Key Differentiator from Lambda Layer Attacks

This attack stores only a benign class name string in the model file. The malicious payload resides in the registered Python class, which is invisible to file-level static analysis. ModelScan returns zero detections because the model file contains no bytecode, serialized functions, or code execution primitives.

Affected Versions

• Format: .keras (Keras v3 ZIP-based format)
• Keras: >= 3.0.0 (all versions with custom layer registry)
• Backend: Agnostic (TensorFlow, JAX, PyTorch)

Reproduction

Prerequisites

• Python 3.9+
• Linux, macOS, or Windows

Steps

1. Clone the repository:

git clone https://huggingface.co/Sanaullah1337/keras-custom-layer-ace-poc
cd keras-custom-layer-ace-poc

2. Install dependencies:

pip install -r requirements.txt

3. Run the proof of concept:

python exploit.py

The script registers the custom layer, loads the model (triggering code execution during deserialization), and verifies the result in a single process.

Expected Output

============================================================
Keras Custom Layer Deserialization — PoC
============================================================
Platform: <Windows|Linux|Darwin>
Output file: <temp-directory>/keras_pwned.txt

[1/3] Custom layer registered: CustomPreprocessingLayer
[2/3] Loading model: demo_model.keras
      Model loaded: sentiment_classifier
      Layers: ['input_features', 'preprocessing', 'hidden_1', 'dropout', 'hidden_2', 'output']

[3/3] Verifying code execution:
      Output file: <temp-directory>/keras_pwned.txt
      Content: ACE_CONFIRMED|host=<hostname>|user=<username>|os=<platform>

Result: Arbitrary code execution confirmed on <platform>

ModelScan Evasion

The model file's config.json was tested against all standard detection patterns:

Pattern	Result
Lambda	Not detected
__reduce__	Not detected
exec(	Not detected
eval(	Not detected
os.system	Not detected
subprocess	Not detected
pickle	Not detected
marshal	Not detected
base64	Not detected
__import__	Not detected

Result: 0/10 detection patterns triggered. The model file contains only the string "CustomPreprocessingLayer" as the link to executable code.

Attack Scenarios

Supply Chain via PyPI

An attacker publishes a package containing a registered Keras custom layer. The package appears to provide legitimate ML preprocessing utilities. Models distributed with this package appear benign to all static scanners. When a victim installs the package and loads the model, arbitrary code executes with the victim's privileges.

Dependency Confusion

An organization uses an internal package that registers Keras custom layers. An attacker publishes a package with the same name on a public registry. The malicious package registers classes with identical names but hidden payloads. CI/CD pipelines that load models will trigger code execution, potentially leaking credentials and source code.

HuggingFace Model Distribution

A model uploaded to HuggingFace references a custom layer class. The model passes all automated security scans. The companion package is presented as an optional dependency for enhanced performance. Users who install the package and load the model are compromised.

Impact

Arbitrary Python code execution on keras.models.load_model(). The victim only needs to load a model file. The attack evades all current static analysis tools. Any Keras v3 user regardless of backend is affected.

Mitigation

• Load untrusted models in sandboxed environments (container, VM, restricted subprocess)
• Audit all registered Keras serializables in the Python environment before loading models with unknown custom layers
• Prefer safetensors format for model weight storage when custom layers are not required
• Keras should implement a safe loading mode that restricts deserialization to built-in layers only

Repository Contents

File	Purpose
demo_model.keras	Malicious Keras v3 model file
exploit.py	Registers the custom layer class
requirements.txt	Python dependencies
README.md	This documentation