Security scanner for AI/ML model files
Project description
TensorTrap
This is a novel Security scanner for AI/ML model files. It detects malicious code in pickle, safetensors, and GGUF files before loading them into workflows. It also checks output files to see if the model files generated malicious code embedded within media files (e.g., jpeg, png, mp4) that could harm your environment when opening/viewing.
Why TensorTrap?
AI model files can contain executable code. Pickle files in particular can run arbitrary Python when loaded. TensorTrap analyzes model files without executing them, identifying dangerous patterns before they can harm your system.
Key statistics:
- 83.5% of Hugging Face models use pickle-based formats (arbitrary code execution risk)
- 2.1 billion monthly downloads from Hugging Face alone
- 100+ confirmed malicious models discovered on public repositories
Platform Support
TensorTrap is cross-platform and runs on all major operating systems:
| Platform | Status | CI Tested |
|---|---|---|
| Linux | Full Support | Ubuntu (Python 3.10-3.12) |
| Windows | Full Support | Windows Server (Python 3.10-3.12) |
| macOS | Full Support | macOS (Python 3.10-3.12) |
All core functionality works identically across platforms. TensorTrap uses pure Python with cross-platform libraries (pathlib, struct, zipfile), ensuring consistent behavior everywhere.
Installation
Windows (Recommended: Standalone Executable)
No Python installation required. Download and run:
- Go to the Releases page
- Download
tensortrap-windows-x64.exe - Move it to a folder in your PATH (e.g.,
C:\Program Files\TensorTrap\) - Open Command Prompt or PowerShell and run:
tensortrap scan .\models\
Tip: To add TensorTrap to your PATH, open System Properties > Environment Variables > edit the
Pathvariable and add the folder where you saved the executable.
Linux / macOS (pip)
pip install tensortrap
Web Dashboard (All Platforms)
The web dashboard provides a browser-based UI for scanning, viewing reports, and managing configuration. Install the web extras:
pip install tensortrap[web]
Development
pip install tensortrap[dev,web]
Web Dashboard
TensorTrap includes a browser-based dashboard that makes scanning and report management accessible without the command line.
Starting the Dashboard
tensortrap serve
This starts a local web server and automatically opens the dashboard in your browser at http://127.0.0.1:7780. To start without opening the browser:
tensortrap serve --no-browser
tensortrap serve --port 8080 # Custom port
Running a Scan
- Click Scan in the left sidebar
- Click Browse to open the folder picker and navigate to the directory you want to scan, or type the path directly
- Adjust scan options if needed (recursive scanning, context analysis, confidence threshold)
- Click Start Scan
- Watch the real-time progress bar as files are scanned
- When complete, click View Full Report to see detailed results
You can navigate to other tabs while a scan is running — the progress is preserved and a banner will show the scan status on other pages.
Viewing Reports
Click Reports in the left sidebar to see all scan reports sorted by date. Click any report to view the full details including:
- Summary statistics (safe files, files with issues, severity breakdown)
- Detailed findings for each flagged file with severity badges
- Confidence scores and recommended actions
- File format, size, and scan time for each result
What To Do With Report Results
- Critical / High severity findings: Do not load these files. Delete them or quarantine them immediately. These indicate known malicious patterns like
os.systemcalls or dangerous pickle opcodes. - Medium severity findings: Investigate further. These may be legitimate patterns (like standard pickle REDUCE opcodes) or potential threats. Check the confidence score — high confidence means the finding is more likely to be a real threat.
- Low / Info findings: Generally informational. Review if you want to be thorough, but these are unlikely to be threats.
- Safe files: No action needed. These files passed all security checks.
Configuration
Click Configuration in the left sidebar to manage all settings from the browser:
Reports
- Report Directory: Where scan reports are saved (use Browse to select a folder)
- Retention: Number of days to keep reports (default: 30, set to 0 to keep forever)
- Report Formats: Choose which formats to generate (HTML, TXT, JSON, CSV)
Web UI
- Port: The port the dashboard runs on (default: 7780)
- Auto-open browser: Whether to open the browser automatically when starting the dashboard
Scheduled Scans
- Enable daily scan: Toggle automatic daily scanning
- Scan Time: What time of day to run the scan (24-hour format, default: 03:00)
- Scan Paths: Directories to scan automatically (one per line)
- Scan Options: Recursive scanning, context analysis, confidence threshold
Click Save Configuration to apply changes, Discard Changes to revert unsaved edits, or Reset to Defaults to restore all settings to their original values.
Running as a Background Service
To have TensorTrap start automatically when you log in:
tensortrap service install # Install and start the service
tensortrap service status # Check if it's running
tensortrap service restart # Restart after config changes
tensortrap service uninstall # Remove the service
Once installed, the dashboard is always available at http://127.0.0.1:7780 — bookmark this URL for easy access.
Note: Background service uses systemd on Linux and launchd on macOS. Logs on macOS are saved to
~/Library/Logs/TensorTrap/.
CLI Usage
Scan a single file:
tensortrap scan model.safetensors
Scan a directory:
tensortrap scan ./models/
Output as JSON (for tooling integration):
tensortrap scan model.pkl --json
Show file info without full scan:
tensortrap info model.safetensors
CLI Options
tensortrap scan <path> [OPTIONS]
Options:
-r, --recursive / -R, --no-recursive Scan directories recursively (default: recursive)
-j, --json Output results as JSON to console
-v, --verbose Show detailed output including info-level findings
--no-hash Skip computing file hashes
--report / --no-report Generate report files (default: enabled for directories)
-o, --report-dir PATH Directory to save reports (overrides config)
-f, --report-formats TEXT Comma-separated formats: txt,json,html,csv (overrides config)
--retain-days INT Days to keep old reports (overrides config, 0 = keep forever)
--context-analysis / --no-context-analysis Context analysis for confidence scoring (default: enabled)
--external-validation Run external tool validation (exiftool/binwalk)
-c, --confidence-threshold FLOAT Minimum confidence to report (0.0-1.0, default: 0.5)
--entropy-threshold FLOAT Entropy threshold for compressed data (0.0-8.0, default: 7.0)
CLI Configuration
TensorTrap stores configuration in ~/.config/tensortrap/config.toml. Manage it from the command line:
tensortrap config init # Interactive setup
tensortrap config show # Display current settings
tensortrap config set <key> <value> # Update a setting
tensortrap config reset # Restore defaults
Report Generation
When scanning directories, TensorTrap automatically generates reports:
# Scan with configured report formats (default)
tensortrap scan ./models/
# Disable report generation
tensortrap scan ./models/ --no-report
# Specific formats only
tensortrap scan ./models/ -f txt,html
# Custom output directory
tensortrap scan ./models/ -o ./reports/
Reports are saved with timestamps: tensortrap_report_YYYYMMDD_HHMMSS.{txt,json,html,csv}
Supported Formats
| Format | Extensions | Risk Level |
|---|---|---|
| Pickle | .pkl, .pickle, .pt, .pth, .bin, .ckpt, .joblib | High (code execution) |
| PyTorch ZIP | .pt, .pth (ZIP archives) | High (internal pickles) |
| Safetensors | .safetensors | Low (data only) |
| GGUF | .gguf | Medium (template injection) |
| ONNX | .onnx | Medium (path traversal) |
| Keras/HDF5 | .h5, .hdf5, .keras | High (Lambda layers, pickle) |
| YAML | .yaml, .yml | Medium (unsafe deserialization) |
| ComfyUI | .json (workflows) | High (eval nodes) |
| Images | .png, .jpg, .gif, .svg, .webp, .bmp, .tiff, .ico | Medium (polyglot attacks) |
| Video | .mp4, .mkv, .avi, .mov, .webm, .flv, .wmv | Medium (polyglot attacks) |
What We Detect
Pickle Files
- Dangerous imports: os, subprocess, socket, builtins, sys, etc.
- Code execution opcodes: REDUCE, BUILD, GLOBAL, INST, NEWOBJ
- Known malicious patterns: os.system, subprocess.Popen, eval, exec
- Nested pickle attacks: pickle importing pickle
Safetensors Files
- Oversized headers: Potential DoS attacks
- Embedded payloads: Pickle data hidden in metadata
- Suspicious patterns: Code snippets in metadata
- Invalid structure: Malformed headers, bad tensor offsets
GGUF Files
- Invalid format: Wrong magic number, unknown versions
- Jinja template injection: CVE-2024-34359 patterns
- Anomalous structure: Excessive tensor/metadata counts
- Suspicious metadata: Code patterns in metadata values
ONNX Files
- Path traversal: CVE-2024-27318, CVE-2024-5187 via external_data
- Suspicious external references: Access to system files
- Arbitrary file read/write: Via malicious external data paths
Keras/HDF5 Files
- Lambda layers: Arbitrary code execution on load
- Embedded pickle: Pickle-serialized custom objects
- Suspicious config patterns: eval(), exec(), os.system()
YAML Configuration Files
- Unsafe deserialization: !!python/object tags (CVE-2025-50460)
- Code execution: subprocess, os.system patterns
- Dynamic imports: import patterns
ComfyUI Workflows
- Vulnerable nodes: ACE_ExpressionEval, HueAdjust (CVE-2024-21576/77)
- Code execution: eval() patterns in node inputs
- Arbitrary code: Malicious workflow structures
Polyglot & Media Files (Defense-in-Depth)
- Extension mismatch: Pickle/archive disguised as image (CVE-2025-1889)
- Archive-in-image: ZIP/7z/RAR appended to valid images
- Archive-in-video: Archives appended to video files
- SVG script injection: JavaScript, onclick handlers, data URIs
- Metadata payloads: Malicious code in EXIF/XMP metadata
- Double extensions: Tricks like
model.pkl.png - Trailing data: Hidden data after image end markers
- MKV attachments: Embedded files in Matroska containers
Additional Detections
- Magic byte analysis: Detects disguised pickle files (CVE-2025-1889)
- 7z archives: nullifAI bypass detection (CVE-2025-1716)
- Obfuscation: Base64, hex encoding, compression, high entropy
- PyTorch archives: Extracts and scans internal pickle files
Benchmark Results
TensorTrap achieves 100% detection rate on comprehensive security benchmarks with zero false positives.
Overall Results
| Metric | Result |
|---|---|
| Overall Accuracy | 43/43 (100%) |
| Malicious Detected | 34/34 (100%) |
| False Positives | 0 |
| False Negatives | 0 |
Detection by Category
| Category | Detection Rate |
|---|---|
| Pickle Bypass | 9/9 (100%) |
| JFrog Zero-Days | 6/6 (100%) |
| Polyglot Attacks | 4/4 (100%) |
| GGUF (Jinja Injection) | 1/1 (100%) |
| ONNX (Path Traversal) | 2/2 (100%) |
| YAML (Unsafe Deserialization) | 2/2 (100%) |
| ComfyUI (ACE/Eval) | 2/2 (100%) |
| Keras/HDF5 (Lambda Layer) | 2/2 (100%) |
| Safetensors | 3/3 (100%) |
| SVG (Script Injection) | 3/3 (100%) |
| Benign (No FP) | 9/9 (100%) |
CVE Coverage
| CVE | Description | Detection |
|---|---|---|
| CVE-2025-1716 | nullifAI 7z/pip bypass | 2/2 (100%) |
| CVE-2025-1889 | ZIP trailing data bypass | 2/2 (100%) |
| CVE-2025-10155 | Extension bypass (.bin/.pt) | 2/2 (100%) |
| CVE-2025-10156 | ZIP zeroed CRC bypass | 1/1 (100%) |
| CVE-2025-10157 | asyncio/_posixsubprocess bypass | 3/3 (100%) |
| CVE-2024-34359 | GGUF Jinja template injection | 1/1 (100%) |
| CVE-2024-27318 | ONNX path traversal | 1/1 (100%) |
| CVE-2024-5187 | ONNX arbitrary file read | 1/1 (100%) |
| CVE-2025-50460 | YAML unsafe deserialization | 1/1 (100%) |
| CVE-2024-21576 | ComfyUI ACE eval | 1/1 (100%) |
| CVE-2024-21577 | ComfyUI HueAdjust eval | 1/1 (100%) |
Running Benchmarks
# Run comprehensive benchmark suite
python tests/benchmark_comprehensive.py --all
# Setup only (generate test files)
python tests/benchmark_comprehensive.py --setup
# Run tests only (after setup)
python tests/benchmark_comprehensive.py --run
# View latest report
python tests/benchmark_comprehensive.py --report
Exit Codes
0: All files safe (no critical/high findings)1: Threats detected (critical or high severity findings)
Example Output
Collecting files from ./models/...
Found 15 model file(s)
Scanning: model.pkl ━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 15/15 0:00:02
model.pkl (pickle) - THREATS DETECTED
Severity Finding Action
!! CRITICAL Known malicious call: os.system DO NOT LOAD. Delete this file immediately.
* MEDIUM REDUCE opcode found 1 time(s) Normal for pickle models. Convert to safetensors.
Scanned 15 file(s): 14 safe, 1 with issues
1 critical, 1 medium
Reports saved:
TXT: ./tensortrap_report_20251211_120000.txt
JSON: ./tensortrap_report_20251211_120000.json
HTML: ./tensortrap_report_20251211_120000.html
CSV: ./tensortrap_report_20251211_120000.csv
JSON Output
{
"report_type": "tensortrap_security_scan",
"scan_target": "./models/",
"scan_date": "2025-12-11T12:00:00",
"summary": {
"total_files": 1,
"safe_files": 0,
"files_with_issues": 1,
"findings_by_severity": {"critical": 1, "medium": 1}
},
"results": [
{
"filepath": "model.pkl",
"format": "pickle",
"is_safe": false,
"max_severity": "critical",
"findings": [
{
"severity": "critical",
"message": "Known malicious call: os.system",
"location": 0,
"details": {"module": "os", "function": "system"},
"recommendation": "DO NOT LOAD. Delete this file immediately."
}
],
"scan_time_ms": 1.23,
"file_size": 256,
"file_hash": "abc123..."
}
]
}
Defense in Depth
TensorTrap focuses on AI model file security. For comprehensive protection of your AI workflow, we recommend combining TensorTrap with these complementary tools:
Recommended Security Stack
| Tool | Purpose | Install |
|---|---|---|
| TensorTrap | AI model file scanning | pip install tensortrap |
| Stego | Steganography detection | See stego-toolkit |
| YARA | Pattern-based malware detection | apt install yara / yara.readthedocs.io |
| RKHunter | Rootkit detection | apt install rkhunter |
| ClamAV | General antivirus | apt install clamav |
What Each Tool Catches
AI Workflow Security
─────────────────────────────────────────────────────────────
Downloaded Models Generated Output System Level
───────────────── ──────────────── ────────────
TensorTrap Stego RKHunter
- Pickle exploits - Hidden data - Rootkits
- Format attacks - Steganography - Backdoors
- Polyglot files
ClamAV
YARA - Known malware
- Known signatures - Viruses
Quick Setup
Linux (pip):
pip install tensortrap # CLI only
pip install tensortrap[web] # CLI + web dashboard
# Optional: full security stack
sudo apt update
sudo apt install yara rkhunter clamav clamav-daemon
sudo freshclam
Windows (Standalone Executable):
Download tensortrap-windows-x64.exe from the Releases page. No Python required.
# Scan models
tensortrap scan .\models\
tensortrap scan $env:USERPROFILE\Downloads\*.pt
Windows (pip):
pip install tensortrap
pip install tensortrap[web] # For the web dashboard
macOS (pip):
pip install tensortrap
pip install tensortrap[web] # For the web dashboard
# Optional: Install YARA via Homebrew
brew install yara
macOS / Linux (Standalone Executable):
Pre-built binaries are also available on the Releases page:
tensortrap-linux-x64tensortrap-macos-arm64(Apple Silicon)tensortrap-macos-x64(Intel)
Read More at M2Dynamics.us
[https://m2dynamics.us/2026/01/11/tensortrap/]
Contributing
Contributions welcome! See CONTRIBUTING.md for guidelines.
# Clone the repo
git clone https://github.com/realmarauder/TensorTrap.git
cd TensorTrap
# Install dev dependencies
pip install -e ".[dev,web]"
# Run tests
pytest
# Run linting
ruff check src/
mypy src/
License
MIT License - see LICENSE.
About
TensorTrap is developed by M2 Dynamics, specializing in AI/ML security consulting.
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Filter files by name, interpreter, ABI, and platform.
If you're not sure about the file name format, learn more about wheel file names.
Copy a direct link to the current filters
File details
Details for the file tensortrap-1.2.0.tar.gz.
File metadata
- Download URL: tensortrap-1.2.0.tar.gz
- Upload date:
- Size: 100.0 kB
- Tags: Source
- Uploaded using Trusted Publishing? Yes
- Uploaded via: twine/6.1.0 CPython/3.13.7
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
c5e77942d9b8b7629109d260c89d5d001a5babcdd250c39c4bbb7c756149b460
|
|
| MD5 |
65f92e2fed5ae3387fa5b9060c0de3d9
|
|
| BLAKE2b-256 |
e532c95a966d0694ac8134412e97b253ea751ce46e35c8e8be0b9feac25e00f7
|
Provenance
The following attestation bundles were made for tensortrap-1.2.0.tar.gz:
Publisher:
publish-pypi.yml on realmarauder/TensorTrap
-
Statement:
-
Statement type:
https://in-toto.io/Statement/v1 -
Predicate type:
https://docs.pypi.org/attestations/publish/v1 -
Subject name:
tensortrap-1.2.0.tar.gz -
Subject digest:
c5e77942d9b8b7629109d260c89d5d001a5babcdd250c39c4bbb7c756149b460 - Sigstore transparency entry: 1190004582
- Sigstore integration time:
-
Permalink:
realmarauder/TensorTrap@f24ae80747e799566b5cd270e0538232a9af87ad -
Branch / Tag:
refs/tags/v1.2.0 - Owner: https://github.com/realmarauder
-
Access:
public
-
Token Issuer:
https://token.actions.githubusercontent.com -
Runner Environment:
github-hosted -
Publication workflow:
publish-pypi.yml@f24ae80747e799566b5cd270e0538232a9af87ad -
Trigger Event:
push
-
Statement type:
File details
Details for the file tensortrap-1.2.0-py3-none-any.whl.
File metadata
- Download URL: tensortrap-1.2.0-py3-none-any.whl
- Upload date:
- Size: 119.2 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? Yes
- Uploaded via: twine/6.1.0 CPython/3.13.7
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
b9f085838417b7014047f47d122f8fea453bbb5b598a0e2dc823fdb7e69736ad
|
|
| MD5 |
53b111bc79a85cc2c6ad675c366a9d67
|
|
| BLAKE2b-256 |
13353e6978d61a4ad487a4860e86fcab04573694aa5c8000d09d0912ecadf5e8
|
Provenance
The following attestation bundles were made for tensortrap-1.2.0-py3-none-any.whl:
Publisher:
publish-pypi.yml on realmarauder/TensorTrap
-
Statement:
-
Statement type:
https://in-toto.io/Statement/v1 -
Predicate type:
https://docs.pypi.org/attestations/publish/v1 -
Subject name:
tensortrap-1.2.0-py3-none-any.whl -
Subject digest:
b9f085838417b7014047f47d122f8fea453bbb5b598a0e2dc823fdb7e69736ad - Sigstore transparency entry: 1190004586
- Sigstore integration time:
-
Permalink:
realmarauder/TensorTrap@f24ae80747e799566b5cd270e0538232a9af87ad -
Branch / Tag:
refs/tags/v1.2.0 - Owner: https://github.com/realmarauder
-
Access:
public
-
Token Issuer:
https://token.actions.githubusercontent.com -
Runner Environment:
github-hosted -
Publication workflow:
publish-pypi.yml@f24ae80747e799566b5cd270e0538232a9af87ad -
Trigger Event:
push
-
Statement type: