Security scanner detecting Python Pickle files performing suspicious actions.
Getting started
Scan a malicious model on Hugging Face:
pip install picklescan
picklescan --huggingface ykilcher/totally-harmless-model
The scanner reports that the Pickle is calling eval() to execute arbitrary code:
https://huggingface.co/ykilcher/totally-harmless-model/resolve/main/pytorch_model.bin:archive/data.pkl: global import '__builtin__ eval' FOUND
----------- SCAN SUMMARY -----------
Scanned files: 1
Infected files: 1
Dangerous globals: 1
The scanner can also load Pickles from local files, directories, URLs, and zip archives (a-la PyTorch):
picklescan --path downloads/pytorch_model.bin
picklescan --path downloads
picklescan --url https://huggingface.co/sshleifer/tiny-distilbert-base-cased-distilled-squad/resolve/main/pytorch_model.bin
To scan Numpy's .npy files, pip install the numpy package first.
Usage
Exit codes
The scanner exit status codes are (a-la ClamAV):
0: scan did not find malware1: scan found malware2: scan failed
Filtering files and directories
When scanning directories, files and subdirectories can be filtered using regular expressions (again modeled after ClamAV). Each option can be specified multiple times:
| Option | Description |
|---|---|
--exclude=REGEX |
Don't scan files whose path matches the regex |
--include=REGEX |
Only scan files whose path matches the regex |
--exclude-dir=REGEX |
Don't descend into directories whose path matches the regex |
--include-dir=REGEX |
Only descend into directories whose path matches the regex |
Key behaviors:
- Excludes always win over includes. A file or directory matching both an exclude and an include pattern is skipped.
- Multiple patterns OR together. A file is included if it matches any
--includepattern. - No includes = everything eligible. Include patterns only narrow the scan when specified.
--exclude-dirprunes traversal. The directory and all of its contents are skipped entirely.
# Only scan .pkl files, skip the cache/ subdirectory
picklescan --path models/ --include='\.pkl$' --exclude-dir='cache'
Develop
Create and activate the conda environment (miniconda is sufficient):
conda env create -f conda.yaml
conda activate picklescan
Install the package in editable mode to develop and test:
python3 -m pip install -e .
Edit with VS Code:
code .
Run unit tests:
pytest tests
Run manual tests:
- Local PyTorch (zip) file
mkdir downloads
wget -O downloads/pytorch_model.bin https://huggingface.co/ykilcher/totally-harmless-model/resolve/main/pytorch_model.bin
picklescan -l DEBUG -p downloads/pytorch_model.bin
- Remote PyTorch (zip) URL
picklescan -l DEBUG -u https://huggingface.co/prajjwal1/bert-tiny/resolve/main/pytorch_model.bin
Lint the code:
black src tests --line-length 140
flake8 src tests --count --show-source
Publish the package to PyPI: bump the package version in setup.cfg and create a GitHub release. This triggers the publish workflow.
Alternative manual steps to publish the package:
python3 -m pip install --upgrade pip
python3 -m pip install --upgrade build
python3 -m build
python3 -m twine upload dist/*
Test the package: bump the version of picklescan in conda.test.yaml and run
conda env remove -n picklescan-test
conda env create -f conda.test.yaml
conda activate picklescan-test
picklescan --huggingface ykilcher/totally-harmless-model
Tested on Linux 5.10.102.1-microsoft-standard-WSL2 x86_64 (WSL2).
References
- pickledoc -- Non-official but in-depth documentation of the Pickle file format
- pickledbg -- Step-by-step Pickle dissassembly debugger
- pickletools.py -- The official "documentation" of the Pickle file format (where documentation == code).
- Machine Learning Attack Series: Backdooring Pickle Files, Johann Rehberger, 2022
- Hugging Face Pickle Scanning, Luc Georges, 2022
- The hidden dangers of loading open-source AI models (ARBITRARY CODE EXPLOIT!, Yannic Kilcher, 2022
- Secure Machine Learning at Scale with MLSecOps, Alejandro Saucedo, 2022
- Backdooring Pickles: A decade only made things worse, ColdwaterQ, DEFCON 2022
- Never a dill moment: Exploiting machine learning pickle files, Evan Sultanik, 2021 (tool: Fickling)
- Exploiting Python pickles, David Hamann, 2020
- Dangerous Pickles - malicious python serialization, Evan Sangaline, 2017
- Python Pickle Security Problems and Solutions, Travis Cunningham, 2015
- Arbitrary code execution with Python pickles, Stephen Checkoway, 2013
- Sour Pickles, A serialised exploitation guide in one part, Marco Slaviero, BlackHat USA 2011 (see also: doc, slides)
1 maintainer
Metadata
1 maintainer
Release files for picklescan 1.0.5
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Source distribution (sdist)
| File | Size | Uploaded | |
|---|---|---|---|
| picklescan-1.0.5.tar.gz | 31.3 kB | Details |
Built distribution (wheel)
| File | Interpreter | ABI | Platform | Reset |
|---|---|---|---|---|
| picklescan-1.0.5-py3-none-any.whl | Python 3 | none | any | Details |
Total release size: 55.4 kB
Release files / picklescan-1.0.5.tar.gz
| Download URL | picklescan-1.0.5.tar.gz |
|---|---|
| Size | 31.3 kB |
| Tags | Source |
|
SHA-256 checksum How to use checksums |
1e220fc0e7eb2021eeb59ee71ce22d6a889ffb9f4d86d6f8272b8defe83ab9dd
|
|
BLAKE2b-256 checksum How to use checksums |
7c19992abcb598424f3d37f3f50783625755c6f49af884cb0fe1365f34c619bf
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/6.2.0 CPython/3.9.25
|
Release files / picklescan-1.0.5-py3-none-any.whl
| Download URL | picklescan-1.0.5-py3-none-any.whl |
|---|---|
| Size | 24.1 kB |
| Tags | Python 3 |
|
SHA-256 checksum How to use checksums |
2abb3d4959a7ee9ccb1e1296db7c56f624bf6d345b93a0a9cf4185bc650dc2fd
|
|
BLAKE2b-256 checksum How to use checksums |
43713035a2ca54845ab1990635ce4ffb115ddd7ed3529dd5afec4e22846d546e
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/6.2.0 CPython/3.9.25
|
