Metadata-Version: 2.4
Name: weightguard
Version: 0.1.2
Summary: Multi-format security scanner for ML model artifacts: pickle/PyTorch, Keras Lambda-layer, ONNX custom-op, and GGUF header checks.
Author: Sanjeev Jaiswal
Author-email: Sanjeev Jaiswal <jassics@gmail.com>
License-Expression: MIT
Requires-Dist: typer>=0.12
Requires-Dist: rich>=13.7
Requires-Dist: pydantic>=2.7
Requires-Dist: fickling>=0.1.12
Requires-Dist: safetensors>=0.8.0
Requires-Dist: huggingface-hub>=1.28.0
Requires-Dist: gguf>=0.19.0
Requires-Dist: h5py>=3.16.0
Requires-Dist: onnx>=1.22.0
Requires-Python: >=3.11
Project-URL: Homepage, https://github.com/jassics/weightguard
Project-URL: Repository, https://github.com/jassics/weightguard
Description-Content-Type: text/markdown

# weightguard

Security scanner for ML model artifacts. Point it at a Hugging Face repo, a git
URL, or a local path, and it flags known artifact-level risks — malicious
pickle/PyTorch payloads, Keras `Lambda`-layer code injection, ONNX custom-op
RCE surface, and anomalous/crafted GGUF headers — with a severity and a
concrete mitigation for each finding.

Model files are not inert data. Several common serialization formats can
embed code that executes the moment the file is *loaded*, before any
inference happens. `weightguard` is a static, pre-load check you run before
trusting a downloaded model.

## Install

```bash
pip install weightguard
```

## Usage

```bash
# Hugging Face repo
weightguard scan https://huggingface.co/<org>/<repo>

# git repo
weightguard scan https://github.com/<org>/<repo>.git

# local path
weightguard scan /path/to/model

# CI gate — exit non-zero only above a severity threshold (default: HIGH)
weightguard scan <target> --fail-on CRITICAL
```

Exit codes: `0` clean, `1` a finding at/above `--fail-on`, `2` target could
not be resolved.

## Use as a library

The CLI is a thin wrapper over the same public API — import it directly to
scan programmatically (CI scripts, pre-deploy hooks, MLOps pipelines):

```python
from weightguard import scan, Severity, UnresolvableTarget

try:
    report = scan("https://huggingface.co/<org>/<repo>")  # or a git URL / local path
except UnresolvableTarget as exc:
    raise SystemExit(f"could not resolve target: {exc}")

for finding in report.findings:
    print(finding.severity, finding.detector, finding.title, finding.file)

if report.fails(Severity.HIGH):
    raise SystemExit("blocking: high-severity finding in model artifact")
```

Lower-level pieces are also exported if you want to resolve and scan
separately, or scan a `pathlib.Path` you already have on disk:

```python
from weightguard import resolve, scan_path

path = resolve("https://github.com/<org>/<repo>.git")  # downloads, returns local Path
report = scan_path(path)
```

`scan_path` never executes or deserializes the target files — it's pure
static analysis, safe to run against untrusted artifacts.

## What it checks today

| Format | Detector | Technique |
|---|---|---|
| Pickle / PyTorch (`.pkl`, `.bin`, `.pt`, `.pth`) | Fickling AST analysis | Detects arbitrary-code-execution opcode chains; resistant to the malformed-opcode-stream evasion that defeats denylist scanners |
| SafeTensors | Format check | Flags files that fail to parse as valid SafeTensors (renamed/spoofed files) |
| Keras (`.h5`, `.keras`) | Lambda-layer check | Flags `Lambda` layers, which embed a marshalled Python function executed on load |
| ONNX | Custom-op check | Flags graphs referencing non-standard operator domains (native-code load surface) |
| GGUF | Header/KV anomaly check | Flags malformed headers, oversized KV metadata, unknown quantization types, and tensor sizes inconsistent with the file |

## Development

```bash
uv sync
uv run pytest -q
```

## License

MIT
