AI/ML Attack Surface

AI/ML Attack Surface

Purpose

Detect security vulnerabilities specific to AI/ML pipelines, LLM-backed applications, and data science workflows. These attack surfaces are increasingly common and often overlooked by traditional SAST tools.

When to Use

Activate this skill when reviewing code that:

• Imports ML frameworks (torch, tensorflow, sklearn, transformers, langchain)
• Loads serialized models or data
• Integrates LLM APIs (OpenAI, Anthropic, etc.)
• Processes Jupyter notebooks
• Handles training data pipelines

Vulnerability Categories

1. Unsafe Deserialization in ML Pipelines (CWE-502)

The most critical ML-specific vulnerability. Many ML serialization formats execute arbitrary code on load.

Dangerous Functions:

Framework	Dangerous	Safe Alternative
PyTorch	torch.load(path)	torch.load(path, weights_only=True)
Joblib	joblib.load(path)	Verify source, use safetensors
NumPy	numpy.load(path, allow_pickle=True)	numpy.load(path, allow_pickle=False)
Scikit-learn	joblib.load() / pickle.load()	skops.io with trusted types
TensorFlow	tf.saved_model.load() with custom ops	Verify model provenance
ONNX	Generally safe	Validate graph structure
SafeTensors	Safe by design	Recommended format

Detection:

# PyTorch unsafe load
grep -rn "torch\.load(" --include="*.py" | grep -v "weights_only=True"

# Joblib/sklearn model loading
grep -rn "joblib\.load\|sklearn.*load" --include="*.py"

# NumPy with pickle enabled
grep -rn "numpy\.load\|np\.load" --include="*.py" | grep "allow_pickle"

# Generic unsafe deserialization in ML context
grep -rn "pickle\.load\|pickle\.loads\|dill\.load\|cloudpickle\.load" --include="*.py"

Exploitation: An attacker who can supply a malicious model file achieves arbitrary code execution on the server loading the model. This is especially dangerous in:

• Model registries that accept user uploads
• Transfer learning pipelines pulling models from external sources
• CI/CD pipelines that load models during testing