modelaudit 0.1.4

Static scanning library for detecting malicious code, backdoors, and other security risks in ML model files

ModelAudit

A security scanner for AI models. Quickly check your AIML models for potential security risks before deployment.

Table of Contents

• ModelAudit
  • Table of Contents
  • 🔍 What It Does
  • 🚀 Quick Start
    • Installation
    • Basic Usage
  • ✨ Features
    • Core Capabilities
    • Reporting & Integration
    • Security Detection
  • 🛡️ Supported Model Formats
    • Weight Analysis
  • ⚙️ Advanced Usage
    • Command Line Options
    • Exit Codes
  • 📋 JSON Output Format
  • 🔄 CI/CD Integration
    • Basic Integration
    • Platform Examples
  • 🔧 Troubleshooting
    • Common Issues
  • ⚠️ Limitations
  • 📝 License

🔍 What It Does

ModelAudit scans ML model files for:

• Malicious code execution (e.g., os.system calls in pickled models)
• Suspicious TensorFlow operations (PyFunc, file I/O operations)
• Potentially unsafe Keras Lambda layers with arbitrary code execution
• Dangerous pickle opcodes (REDUCE, INST, OBJ, STACK_GLOBAL)
• Custom ONNX operators and external data integrity issues
• Encoded payloads and suspicious string patterns
• Risky configurations in model architectures
• Suspicious patterns in model manifests and configuration files
• Models with blacklisted names or content patterns
• Malicious content in ZIP archives including nested archives and zip bombs
• Container-delivered models in OCI/Docker layers and manifest files
• GGUF/GGML file integrity and tensor alignment validation
• Anomalous weight patterns that may indicate trojaned models (statistical analysis)
• Enhanced joblib/dill security (format validation, compression bombs, embedded pickle analysis, bypass prevention)
• NumPy array integrity issues (malformed headers, dangerous dtypes)

🚀 Quick Start

Installation

ModelAudit is available on PyPI and requires Python 3.9 or higher.

Basic installation:

pip install modelaudit

With optional dependencies for specific model formats:

# For TensorFlow SavedModel scanning
pip install modelaudit[tensorflow]

# For Keras H5 model scanning
pip install modelaudit[h5]

# For PyTorch model scanning
pip install modelaudit[pytorch]

# For ONNX model scanning
pip install modelaudit[onnx]

# For TensorFlow Lite model scanning
pip install modelaudit[tflite]

# For YAML manifest scanning
pip install modelaudit[yaml]

# For SafeTensors model scanning
pip install modelaudit[safetensors]

# For enhanced pickle support (dill serialization with security validation)
pip install modelaudit[dill]

# For Joblib model scanning (includes scikit-learn integration)
pip install modelaudit[joblib]

# For Flax msgpack scanning
pip install modelaudit[flax]

# Install all optional dependencies
pip install modelaudit[all]

Development installation:

git clone https://github.com/promptfoo/modelaudit.git
cd modelaudit

# Using Rye (recommended)
rye sync --features all

# Or using pip
pip install -e .[all]

Docker installation:

# Pull from GitHub Container Registry
docker pull ghcr.io/promptfoo/modelaudit:latest

# Use specific variants
docker pull ghcr.io/promptfoo/modelaudit:latest-full        # All ML frameworks
docker pull ghcr.io/promptfoo/modelaudit:latest-tensorflow  # TensorFlow only

# Run with Docker
docker run --rm -v $(pwd):/data ghcr.io/promptfoo/modelaudit:latest scan /data/model.pkl

Basic Usage

# Scan a single model
modelaudit scan model.pkl

# Scan an ONNX model
modelaudit scan model.onnx

# Scan multiple models (including enhanced dill/joblib support)
modelaudit scan model1.pkl model2.h5 model3.pt llama-model.gguf model4.joblib model5.dill model6.npy flax-checkpoint.msgpack

# Scan a directory
modelaudit scan ./models/

# Export results to JSON
modelaudit scan model.pkl --format json --output results.json

Example output:

$ modelaudit scan suspicious_model.pkl

────────────────────────────────────────────────────────────────────────────────
ModelAudit Security Scanner
Scanning for potential security issues in ML model files
────────────────────────────────────────────────────────────────────────────────
Paths to scan: suspicious_model.pkl
────────────────────────────────────────────────────────────────────────────────

✓ Scanning suspicious_model.pkl

Active Scanner: pickle
Scan completed in 0.02 seconds
Files scanned: 1
Scanned 156 bytes
Issues found: 2 critical, 1 warnings

1. suspicious_model.pkl (pos 28): [CRITICAL] Suspicious module reference found: posix.system
2. suspicious_model.pkl (pos 52): [WARNING] Found REDUCE opcode - potential __reduce__ method execution

────────────────────────────────────────────────────────────────────────────────
✗ Scan completed with findings

✨ Features

Core Capabilities

• Multiple Format Support: PyTorch (.pt, .pth, .bin), TensorFlow (SavedModel, .pb), Keras (.h5, .hdf5, .keras), SafeTensors (.safetensors), GGUF/GGML (.gguf, .ggml), Pickle (.pkl, .pickle, .ckpt), Joblib (.joblib), NumPy (.npy, .npz), PMML (.pmml), ZIP archives (.zip), Manifests (.json, .yaml, .xml, etc.), Flax (.msgpack, .ckpt)
• Automatic Format Detection: Identifies model formats automatically
• Deep Security Analysis: Examines model internals, not just metadata
• Recursive Archive Scanning: Scans contents of ZIP files and nested archives
• Batch Processing: Scan multiple files and directories efficiently
• Configurable Scanning: Set timeouts, file size limits, custom blacklists

Reporting & Integration

• Multiple Output Formats: Human-readable text and machine-readable JSON
• Detailed Reporting: Scan duration, files processed, bytes scanned, issue severity
• Severity Levels: CRITICAL, WARNING, INFO, DEBUG for flexible filtering
• CI/CD Integration: Clear exit codes for automated pipeline integration

Security Detection

• Code Execution: Detects embedded Python code, eval/exec calls, system commands
• Pickle Security: Analyzes dangerous opcodes, suspicious imports, encoded payloads
• Enhanced Dill/Joblib Analysis: ML-aware scanning with format validation and bypass prevention
• Model Integrity: Checks for unexpected files, suspicious configurations
• Archive Security: Automatic Zip-Slip protection against directory traversal, zip bombs, malicious nested files
• Pattern Matching: Custom blacklist patterns for organizational policies

🛡️ Supported Model Formats

ModelAudit provides specialized security scanners for different model formats:

Format	File Extensions	What We Check
Pickle	.pkl, .pickle, .dill, .bin, .pt, .pth, .ckpt	Malicious code execution, dangerous opcodes, suspicious imports
PyTorch Zip	.pt, .pth	Embedded pickle analysis, suspicious files, custom patterns
PyTorch Binary	.bin	Binary tensor data analysis, embedded content
TensorFlow Lite	.tflite	Extreme tensor shapes, custom ops, FlatBuffer integrity
TensorFlow	SavedModel dirs, .pb	Suspicious operations, file I/O, Python execution
Keras	.h5, .hdf5, .keras	Lambda layers, custom objects, dangerous configurations
ONNX	.onnx	Custom operators, external data validation, tensor integrity
SafeTensors	.safetensors	Metadata integrity, tensor validation
Flax	.msgpack	MessagePack integrity, suspicious code pattern detection, decompression bomb prevention, embedded content analysis
GGUF/GGML	.gguf, .ggml	Header validation, tensor integrity, metadata security checks
Joblib	.joblib	File format validation, compression bomb detection, embedded pickle analysis, ML-aware security filtering
NumPy	.npy, .npz	Array integrity, dangerous dtypes, dimension validation
PMML	.pmml	XML well-formedness, external entity checks, suspicious extensions
ZIP Archives	.zip	Recursive content scanning, zip bombs, directory traversal
Manifests	.json, .yaml, .yml, .xml, .toml, .ini, .cfg, .config, .manifest, .model, .metadata	Suspicious keys, credential exposure, blacklisted patterns

Weight Analysis

ModelAudit can detect anomalous weight patterns that may indicate trojaned models using statistical analysis. This feature is disabled by default for large language models to avoid false positives.

⚙️ Advanced Usage

Command Line Options

# Set maximum file size to scan (1GB limit)
modelaudit scan model.pkl --max-file-size 1073741824

# Stop scanning after a total of 5GB has been processed
modelaudit scan models/ --max-total-size 5368709120

# Add custom blacklist patterns
modelaudit scan model.pkl --blacklist "unsafe_model" --blacklist "malicious_net"

# Set scan timeout (5 minutes)
modelaudit scan large_model.pkl --timeout 300

# Verbose output for debugging
modelaudit scan model.pkl --verbose

Exit Codes

ModelAudit uses different exit codes to indicate scan results:

• 0: Success - No security issues found
• 1: Security issues found (scan completed successfully)
• 2: Errors occurred during scanning (e.g., file not found, scan failures)

📋 JSON Output Format

When using --format json, ModelAudit outputs structured results:

{
  "scanner_names": ["pickle"],
  "start_time": 1750168822.481906,
  "bytes_scanned": 74,
  "issues": [
    {
      "message": "Found REDUCE opcode - potential __reduce__ method execution",
      "severity": "warning",
      "location": "evil.pickle (pos 71)",
      "details": {
        "position": 71,
        "opcode": "REDUCE",
        "ml_context_confidence": 0.0
      },
      "timestamp": 1750168822.482304
    },
    {
      "message": "Suspicious module reference found: posix.system",
      "severity": "critical",
      "location": "evil.pickle (pos 28)",
      "details": {
        "module": "posix",
        "function": "system",
        "position": 28,
        "opcode": "STACK_GLOBAL",
        "ml_context_confidence": 0.0
      },
      "timestamp": 1750168822.482378
    }
  ],
  "has_errors": false,
  "files_scanned": 1,
  "duration": 0.0005328655242919922
}

Each issue includes a message, severity level (critical, warning, info, debug), location, and scanner-specific details.

🔄 CI/CD Integration

ModelAudit is designed to integrate seamlessly into CI/CD pipelines with clear exit codes:

• Exit Code 0: No security issues found
• Exit Code 1: Security issues found (fails the build)
• Exit Code 2: Scan errors occurred (fails the build)

Basic Integration

# Install ModelAudit
pip install modelaudit[all]

# Scan models and fail build if issues found
modelaudit scan models/ --format json --output scan-results.json

# Optional: Upload scan-results.json as build artifact

Platform Examples

GitHub Actions:

- name: Scan models
  run: |
    rye run modelaudit scan models/ --format json --output scan-results.json
    if [ $? -eq 1 ]; then
      echo "Security issues found in models!"
      exit 1
    fi

GitLab CI:

model-security-scan:
  script:
    - pip install modelaudit[all]
    - modelaudit scan models/ --format json --output results.json
  artifacts:
    paths: [results.json]

Jenkins:

sh 'pip install modelaudit[all]'
sh 'modelaudit scan models/ --format json --output results.json'

🔧 Troubleshooting

Common Issues

Installation Problems:

# If you get dependency conflicts
pip install --upgrade pip setuptools wheel
pip install modelaudit[all] --no-cache-dir

# Install with Rye (recommended)
rye sync --features all

# Or with pip
pip install -e .[all]

# If optional dependencies fail, install base package first
pip install modelaudit
pip install tensorflow h5py torch pyyaml safetensors onnx joblib  # Add what you need

Large Models:

# Increase file size limit and timeout for large models
modelaudit scan large_model.pt --max-file-size 5000000000 --timeout 600 --max-total-size 10000000000

Testing:

# Run all tests
rye run pytest

# Run with coverage
rye run pytest --cov=modelaudit

# Run specific test categories
rye run pytest tests/test_pickle_scanner.py -v
rye run pytest tests/test_integration.py -v

# Run tests with all optional dependencies
rye sync --features all
rye run pytest

# Run comprehensive migration test (tests everything including Docker)
./test_migration.sh

Debug Mode:

# Enable verbose output for troubleshooting
modelaudit scan model.pkl --verbose

Development Commands:

# Run linting and formatting with Ruff
rye run ruff check modelaudit/          # Check for linting issues
rye run ruff check --fix modelaudit/    # Fix auto-fixable issues
rye run ruff format modelaudit/         # Format code

# Type checking
rye run mypy modelaudit/

# Build package
rye build

# Publish (maintainers only)
rye publish

Getting Help:

• Use --verbose for detailed output
• Use --format json to see all details
• Check file permissions and format support
• Report issues on the promptfoo GitHub repository

⚠️ Limitations

ModelAudit is designed to find obvious security risks in model files, including direct code execution attempts, known dangerous patterns, malicious archive structures, and suspicious configurations.

What it cannot detect:

• Advanced adversarial attacks or subtle weight manipulation
• Heavily encoded/encrypted malicious payloads
• Runtime behavior that only triggers under specific conditions
• Model poisoning through careful data manipulation

Recommendations:

• Use ModelAudit as one layer of your security strategy
• Review flagged issues manually - not all warnings indicate malicious intent
• Combine with other security practices like sandboxed execution and runtime monitoring
• Implement automated scanning in CI/CD pipelines

📝 License

This project is licensed under the MIT License - see the LICENSE file for details.