Model file threat scanning — detect malicious pickle, safetensors, and ONNX files in agent codebases

[JBAhire]

Summary

Detect and flag potentially malicious AI model files in agent codebases. Pickle files can execute arbitrary code on load, and model files from untrusted sources are a growing supply chain attack vector.

Motivation

AI agents often reference or bundle model files. These files can contain:

• Arbitrary code execution — Python pickle files execute code on pickle.load()
• Trojaned models — models with backdoors that activate on specific inputs
• Embedded payloads — malicious code hidden in model metadata or custom layers
• Unsigned models — no verification of model provenance or integrity

g0 scans agent code but doesn't inspect model artifacts in the project directory.

Proposed Implementation

1. Model File Detection

• Identify model files in project: .pkl, .pickle, .pt, .pth, .bin, .safetensors, .onnx, .h5, .keras, .pb, .tflite
• Flag in g0 inventory as model components
• Include model files in AIBOM output

2. Pickle Deserialization Risk

• Critical finding for any .pkl/.pickle file in a project
• Detect pickle.load(), torch.load(), joblib.load() without safety flags
• Flag torch.load() without weights_only=True
• Detect custom __reduce__ methods in Python classes (code execution vector)

3. SafeTensors Validation

• Verify safetensors files are well-formed
• Check for unusual metadata that could indicate tampering
• Prefer safetensors over pickle (informational finding)

4. Model Provenance Checks

• Detect models loaded from untrusted URLs (not HuggingFace, not verified sources)
• Flag models without hash pinning (no SHA-256 verification)
• Detect from_pretrained() calls without revision pinning
• Check for models loaded via HTTP (not HTTPS)

5. New Security Rules

• AA-SC-200 — Pickle file in project directory (critical)
• AA-SC-201 — torch.load() without weights_only=True (high)
• AA-SC-202 — Model loaded from untrusted source (high)
• AA-SC-203 — Model loaded without hash verification (medium)
• AA-SC-204 — Model loaded over HTTP (high)
• AA-SC-205 — Custom __reduce__ in class used with pickle (critical)
• AA-SC-206 — from_pretrained() without revision pin (medium)

Files to Create/Modify

• src/analyzers/model-scan.ts — model file detection and analysis
• src/rules/builtin/supply-chain/model-*.yaml — model security rules
• Update src/analyzers/rules/supply-chain.ts with pickle/model checks
• Update src/cli/commands/inventory.ts to include model files

Acceptance Criteria

• Detect model files in project directories
• Critical finding for pickle files
• Detect unsafe torch.load() and pickle.load()
• Model provenance checks (source, hash, revision)
• 7+ new supply chain rules for model security
• Model files included in inventory output