Workflow:Trailofbits Fickling Safe ML Model Loading
| Knowledge Sources | |
|---|---|
| Domains | ML_Security, Supply_Chain_Security |
| Last Updated | 2026-02-14 13:00 GMT |
Overview
End-to-end process for securing AI/ML model loading pipelines by activating Fickling's import allowlist enforcement on the Python pickle module.
Description
This workflow establishes a secure environment for loading pickle-serialized ML models (PyTorch, scikit-learn, NumPy, etc.) by monkey-patching Python's built-in pickle module at runtime. Fickling replaces the standard pickle unpickler with a restricted variant that checks every import against a curated allowlist of safe ML library imports. Any import not on the allowlist raises an UnsafeFileError, preventing arbitrary code execution from malicious model files.
The protection operates at the pickle.Unpickler level, intercepting find_class calls that resolve module/name pairs during deserialization. This means protection is transparent to downstream code that uses pickle.load(), torch.load(), or numpy.load() under the hood.
Usage
Execute this workflow when deploying ML inference pipelines or training environments that load pre-trained models from untrusted or semi-trusted sources (e.g., HuggingFace Hub, shared model repositories). It is especially relevant when models are serialized using Python pickle format (.pt, .pth, .pkl, .bin files) and when the loading code cannot be modified to use safer alternatives like SafeTensors.
Execution Steps
Step 1: Install Fickling
Install the fickling package with optional PyTorch support if working with PyTorch model files. The package has minimal dependencies and supports Python 3.9 through 3.13.
Key considerations:
- Use the [torch] extra if working with PyTorch model formats
- Fickling itself has very few dependencies to minimize attack surface
Step 2: Activate Safe ML Environment
Call the activation function once at process startup, before any model loading occurs. This globally patches the pickle and _pickle C-extension modules, replacing pickle.load, pickle.loads, and pickle.Unpickler with Fickling's allowlist-enforcing variants.
Key considerations:
- Must be called before any model loading code executes
- The hook is global and affects all pickle operations in the process
- Optionally pass an also_allow list to permit additional imports specific to your use case
Step 3: Load Models Through Standard APIs
Use the standard model loading APIs (torch.load(), pickle.load(), numpy.load(), etc.) as normal. Fickling's hooks transparently intercept deserialization and enforce the allowlist.
Key considerations:
- No changes to existing model loading code are required
- Protection works for any code path that eventually calls pickle.Unpickler
- PyTorch's internal torch.load() calls are intercepted via the Unpickler class replacement
Step 4: Handle Blocked Imports
When a model file contains imports not on the allowlist, Fickling raises an UnsafeFileError with details about the blocked import. Catch this exception to implement graceful error handling or logging.
Key considerations:
- The exception includes the module and function name that was blocked
- Review blocked imports to determine if they should be added to the custom allowlist
- Never blindly add blocked imports without verifying they cannot enable arbitrary code execution
Step 5: Deactivate Protection (Optional)
If needed, restore the original pickle behavior by calling the deactivation function. This restores the original pickle.load, pickle.loads, and pickle.Unpickler references.
Key considerations:
- Only deactivate if you have a specific need to load untrusted pickle data
- After deactivation, all pickle operations revert to standard unsafe behavior