Injecting backdoor into a ML model

I am Harish SG, a security researcher who studies Masters in Cybersecurity at UT Dallas,previously hunted on the Microsoft Bug Bounty Program and Google VRP.

I am sharing this article for security awareness and educational purposes only.

In this article, I explained , how I injected malicious code into .pkl file and how to defend against these kind of pickle based supply chain attack in python environment.

What is pickle file in python?

In Python, a “pickle” file is used to serialize and deserialize objects. Serialization is the process of converting a data structure or object into a format that can be easily stored (for instance, in a file) or transmitted, and later reconstructed. The module that facilitates this process in Python is aptly named pickle

What is pickle file in the context of machine learning?

In the context of machine learning, a pickle file is often used to serialize and save a trained model so that it can be loaded and used later without having to retrain it. This becomes particularly handy when model training is time-consuming or resource-intensive. Once a model is trained, it can be saved (or “pickled”) and then shared or deployed in different environments.

Here’s why and how pickle is used in machine learning:

1. Convenience: After spending hours or even days training a complex model, you don’t want to redo that process every time you want to make a prediction or restart your application. By saving the trained model to a pickle file, you can reload it in seconds and use it for predictions.
2. Sharing & Deployment: If you’re working in a team or want to deploy a model into production, you can train your model, save it to a pickle file, and then share that file. Others can then use the pickled model without needing access to the original training data or going through the training process.
3. Portability: Models can be transported between different platforms using pickle files, provided the platforms support the same Python environment and library versions.

how the pickle module implements its serialization and deserialization processes using a stack-based virtual machine.

Stack-based Virtual Machine: The pickle format is essentially a stack-based instruction set for a virtual machine. This VM reads pickle instructions and then carries out operations based on these instructions.

Instructions: The pickle data consists of various instructions that guide the virtual machine on how to reconstruct the original object. These instructions might include

Pushing an item onto the stack.

• Popping an item from the stack.
• Building a list using the top N items from the stack.
• Creating an instance of a certain class using items from the stack as arguments, and so on.

Execution: When you unpickle data, the pickle virtual machine reads and executes these instructions in sequence. As it processes the instructions, it rebuilds the original data structure. The machine uses a stack to keep track of the data as it’s being reconstructed.

Disassembling .pkl file

1. \x80 PROTO 4: This sets the pickle protocol version to 4. The pickle protocol version determines the format in which objects are serialized and deserialized. Protocol version 4 was added in Python 3.4 and offers support for very large objects, pickling more kinds of objects, and some data format optimizations.
2. \x95 FRAME 25: This specifies the size of the subsequent frame, which is 25 bytes long. The frame-based structure was introduced in protocol 4 to provide better support for large pickled data streams.
3. \x8c SHORT_BINUNICODE 'Backdooring ML.......': This pushes the Unicode string 'Backdooring ML.......' onto the pickle VM's stack. The SHORT_BINUNICODE opcode is used to represent short Unicode strings, where "short" usually refers to strings whose byte representation is less than 256 bytes.
4. \x94 MEMOIZE (as 0): This stores the top stack item in the memo (a kind of cache used in the pickling process). This allows future references in the pickle stream to refer back to this object rather than duplicating its content. The (as 0) indicates that the item is being stored with the key 0 in the memo.
5. . (STOP): This instruction marks the end of the pickle stream.

lets hack into pickle serialized ML model

In the above screenshot I wrote an example python script which injects python code (print(“you ‘ve been pwned”)) into an example .pkl file.

Analyzing backdoored .pkl file

In the above screenshot , I used pickletools.dis function to disassemble malicious .pkl file.

When this malicious pickle is unpickled, it would execute arbitrary code — in this case, a simple print statement: print("you've been pwned !").

Let’s break down what’s happening in the provided payload:

1. \x80 PROTO 4: This is setting the pickle protocol version to 4.
2. \x95 FRAME 19: This specifies the size of the subsequent frame.
3. c GLOBAL 'builtins exec': This instruction tells the pickle machine to fetch the exec function from Python's built-in functions.
4. ( (MARK): This marks the start of a tuple.
5. \x8c SHORT_BINUNICODE 'print("you've been pwned !")': This is pushing the string 'print("you've been pwned !")' onto the pickle VM's stack.
6. t TUPLE (MARK at 26): This ends the tuple started at the mark.
7. R REDUCE: This applies the previously mentioned function (exec) to the tuple, which contains our string. In essence, it executes the command exec('print("you've been pwned !")').
8. \x8c SHORT_BINUNICODE 'Backkdooring ML': This pushes the string 'Backkdooring ML' onto the stack, but it's not actually being used here. It's more for the message and doesn't play a part in the malicious action.
9. \x94 MEMOIZE (as 0): This caches the result for later reuse.
10. . (STOP): This instruction ends the pickle stream.

Note : highlighted code responsible for injecting malicious python code into .pkl file

The instructions that pose a threat are STACK_GLOBAL, GLOBAL and REDUCE.

REDUCE is what tells the unpickler to execute the function with the provided arguments and *GLOBAL instructions are telling the unpickler to import stuff.

Injecting malware into OpenAI FAISS index .pkl file

code to inject backdoor into legitimate .pkl file

Backdoor code got executed while loading legitimate .pkl

From the above screenshots , you can observe that I successfully injected backdoor into OpenAI faiss pkl , this demonstrates process of loading an existing pickle file, modifying it to include arbitrary Python code, and then saving the modified pickle back to disk. The payload inserted is intended to be executed whenever the pickle file is loaded.

Steps to Defend your ML model and IT infrastructure from malicious pickle file

Use file hashing to validate against tampering

Use Alternative Serialization Formats:

• Use safer serialization formats such as JSON (json module), XML, or YAML (with careful usage since certain libraries have had their own vulnerabilities). These formats don't support code execution.
• For JSON, use Python’s built-in json module.
• For more complex Python-specific structures, consider using jsonpickle which is like pickle but produces JSON, which is easier to inspect.

Secure the Environment:

• If you absolutely must unpickle data from an untrusted source, consider using a sandboxed environment. Libraries such as pypy-sandbox or using container solutions like Docker can isolate the environment and mitigate the damage a malicious payload can cause.

Least Privilege Principle:

• Run applications that use pickle with the least privilege necessary. This way, if a malicious pickle does get executed, it can do less harm.

Restrict Global Imports:

1. just use pickle, you can override the find_class method to restrict global imports to safe modules or classes. Here's a basic example:

Try hacking LLM : https://github.com/harishsg993010/DamnVulnerableLLMProject

Hacking into Bard : https://infosecwriteups.com/hacking-google-bard-24f9dfa7b455

Follow me on twitter: https://twitter.com/CoderHarish

Follow me on linkedin :https://www.linkedin.com/in/harish-santhanalakshmi-ganesan-31ba96171/