Metadata-Version: 2.4
Name: modelaudit
Version: 0.1.4
Summary: Static scanning library for detecting malicious code, backdoors, and other security risks in ML model files
Project-URL: Repository, https://github.com/promptfoo/modelaudit
Project-URL: Homepage, https://github.com/promptfoo/modelaudit
Author-email: Ian Webster <ian@promptfoo.dev>
License: MIT
License-File: LICENSE
Keywords: ai,ml,model-scanning,pickle,pytorch,security,tensorflow
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Security
Requires-Python: >=3.9
Requires-Dist: click>=8.1.3
Requires-Dist: defusedxml>=0.7.1
Requires-Dist: numpy>=1.19.0
Requires-Dist: scipy>=1.5.0
Requires-Dist: yaspin>=2.3.0
Provides-Extra: all
Requires-Dist: dill>=0.3.0; extra == 'all'
Requires-Dist: h5py>=3.1; extra == 'all'
Requires-Dist: joblib>=1.0.0; extra == 'all'
Requires-Dist: msgpack>=1.0.0; extra == 'all'
Requires-Dist: onnx>=1.12.0; extra == 'all'
Requires-Dist: pyyaml>=6.0; extra == 'all'
Requires-Dist: safetensors>=0.4.0; extra == 'all'
Requires-Dist: scikit-learn>=1.0.0; extra == 'all'
Requires-Dist: tensorflow>=2.13.0; extra == 'all'
Requires-Dist: tflite>=2.18.0; extra == 'all'
Requires-Dist: torch>=1.6; extra == 'all'
Provides-Extra: dill
Requires-Dist: dill>=0.3.0; extra == 'dill'
Provides-Extra: flax
Requires-Dist: msgpack>=1.0.0; extra == 'flax'
Provides-Extra: h5
Requires-Dist: h5py>=3.1; extra == 'h5'
Provides-Extra: joblib
Requires-Dist: joblib>=1.0.0; extra == 'joblib'
Requires-Dist: scikit-learn>=1.0.0; extra == 'joblib'
Provides-Extra: onnx
Requires-Dist: onnx>=1.12.0; extra == 'onnx'
Provides-Extra: pytorch
Requires-Dist: torch>=1.6; extra == 'pytorch'
Provides-Extra: safetensors
Requires-Dist: safetensors>=0.4.0; extra == 'safetensors'
Provides-Extra: tensorflow
Requires-Dist: tensorflow>=2.13.0; extra == 'tensorflow'
Provides-Extra: tflite
Requires-Dist: tflite>=2.18.0; extra == 'tflite'
Provides-Extra: yaml
Requires-Dist: pyyaml>=6.0; extra == 'yaml'
Description-Content-Type: text/markdown

# ModelAudit

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

[![PyPI version](https://badge.fury.io/py/modelaudit.svg)](https://pypi.org/project/modelaudit/)
[![Python 3.9+](https://img.shields.io/badge/python-3.9+-blue.svg)](https://www.python.org/downloads/release/python-390/)

<img width="989" alt="image" src="https://www.promptfoo.dev/img/docs/modelaudit/modelaudit-result.png" />

## Table of Contents

- [ModelAudit](#modelaudit)
  - [Table of Contents](#table-of-contents)
  - [🔍 What It Does](#-what-it-does)
  - [🚀 Quick Start](#-quick-start)
    - [Installation](#installation)
    - [Basic Usage](#basic-usage)
  - [✨ Features](#-features)
    - [Core Capabilities](#core-capabilities)
    - [Reporting \& Integration](#reporting--integration)
    - [Security Detection](#security-detection)
  - [🛡️ Supported Model Formats](#️-supported-model-formats)
    - [Weight Analysis](#weight-analysis)
  - [⚙️ Advanced Usage](#️-advanced-usage)
    - [Command Line Options](#command-line-options)
    - [Exit Codes](#exit-codes)
  - [📋 JSON Output Format](#-json-output-format)
  - [🔄 CI/CD Integration](#-cicd-integration)
    - [Basic Integration](#basic-integration)
    - [Platform Examples](#platform-examples)
  - [🔧 Troubleshooting](#-troubleshooting)
    - [Common Issues](#common-issues)
  - [⚠️ Limitations](#️-limitations)
  - [📝 License](#-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](https://pypi.org/project/modelaudit/) and requires **Python 3.9 or higher**.

**Basic installation:**

```bash
pip install modelaudit
```

**With optional dependencies for specific model formats:**

```bash
# 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:**

```bash
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:**

```bash
# 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

```bash
# 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:**

```bash
$ 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

```bash
# 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:

```json
{
  "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

```bash
# 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:**

```yaml
- 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:**

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

**Jenkins:**

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

## 🔧 Troubleshooting

### Common Issues

**Installation Problems:**

```bash
# 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:**

```bash
# 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:**

```bash
# 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:**

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

**Development Commands:**

```bash
# 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](https://github.com/promptfoo/promptfoo/issues)

## ⚠️ 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](https://opensource.org/licenses/MIT) - see the [LICENSE](LICENSE) file for details.
