Workflow:Trailofbits Fickling PyTorch Payload Injection
| Knowledge Sources | |
|---|---|
| Domains | Security_Research, Red_Teaming, ML_Security |
| Last Updated | 2026-02-14 13:00 GMT |
Overview
End-to-end process for injecting arbitrary Python code into PyTorch model files to demonstrate pickle deserialization vulnerabilities and test scanner detection capabilities.
Description
This workflow enables security researchers to create proof-of-concept malicious PyTorch model files by injecting code payloads into existing model archives. Fickling provides two injection strategies: insertion mode, which modifies the pickle bytecode within the existing ZIP archive to execute additional code during deserialization; and combination mode, which wraps the original model in a new module whose __reduce__ method triggers the payload while preserving the model's forward pass functionality.
This capability is intended for authorized security testing, red team exercises, evaluating ML supply chain defenses, and benchmarking pickle scanning tools. The injected models can then be used to verify that safety tools (including Fickling itself) correctly detect the malicious payload.
Usage
Execute this workflow when performing authorized penetration testing of ML pipelines, building test datasets for scanner benchmarking, demonstrating pickle deserialization risks to stakeholders, or verifying that security controls detect injected payloads. This workflow requires explicit authorization context.
Execution Steps
Step 1: Prepare a PyTorch Model File
Obtain or create a PyTorch model file in a supported format. The target file should be in PyTorch v1.3 format (ZIP with data.pkl) or TorchScript v1.4 format. Save the model using torch.save() to produce the expected file structure.
Key considerations:
- The model file must be in a format supported by PyTorchModelWrapper (PyTorch v1.3 or TorchScript v1.4)
- Other formats may work with the force=True parameter but results are not guaranteed
- The original model architecture and weights are preserved after injection
Step 2: Wrap with PyTorchModelWrapper
Load the model file into Fickling's PyTorchModelWrapper, which validates the file format, identifies the PyTorch version, and extracts the pickle stream from the ZIP archive.
Key considerations:
- Format validation ensures the file is a recognized PyTorch format
- The wrapper locates data.pkl within the ZIP structure
- TorchScript v1.4 support is experimental and injections may not work with all parsers
Step 3: Define the Payload
Craft the Python code string that will execute during deserialization. The payload is a string of Python code that will be evaluated when the model is loaded.
Key considerations:
- Payloads execute in the context of the deserializing process
- For insertion mode, the payload is injected as additional pickle opcodes using insert_python_exec
- For combination mode, the payload becomes the return value of a __reduce__ method via eval
Step 4: Inject the Payload
Call the injection method with the chosen strategy. Insertion mode modifies the pickle bytecode in-place within the ZIP archive, adding the payload as additional operations. Combination mode creates a new model class that wraps the original and triggers the payload via its __reduce__ method.
Key considerations:
- Insertion mode: modifies data.pkl opcodes; does NOT bypass weights-based unpicklers
- Combination mode: creates an entirely new model wrapper; detectable by type checking
- The overwrite parameter replaces the original file with the injected version
- Output path specifies where the modified model is saved
Step 5: Verify Detection
Load the injected model file through Fickling's safety analysis to confirm the payload is detectable. Use check_safety() or is_likely_safe() to verify the injected file is flagged.
Key considerations:
- A properly functioning scanner should detect the injected payload
- Compare detection results across multiple scanning tools for benchmarking
- Document detection gaps for responsible disclosure or tool improvement