python-security-pitfalls

Python Security Pitfalls (services, data and ML)

Overview

Python's dangerous calls are few and well known — pickle, eval, shell=True, f-string SQL — but data and ML code uses them as everyday tools: models are pickles, configs are YAML, notebooks evaluate whatever they are given. The question is always where the bytes come from.

1. Model and Data Loading (CWE-502) — the data/ML hot spot

Loading a pickle-based file runs code chosen by whoever wrote the file.

# ❌ an uploaded or downloaded model file executes code on load
model = torch.load(path, weights_only=False)
clf = joblib.load(user_upload)
df = pd.read_pickle(s3_key_from_request)
arr = np.load(path, allow_pickle=True)
cfg = yaml.load(text, Loader=yaml.Loader)

# ✅
model.load_state_dict(safetensors.torch.load_file(path))
arr = np.load(path)                    # allow_pickle defaults to False
cfg = yaml.safe_load(text)

• torch.load defaults to weights_only=True from PyTorch 2.6; on an older pinned torch the plain call is unsafe.
• Hugging Face from_pretrained(..., trust_remote_code=True) runs Python from the model repository; with a user-chosen repo id it is remote code execution.
• Keras load_model(..., safe_mode=False) allows Lambda-layer code.
• Celery or Kombu with accept_content=["pickle"], pickled values in a shared Redis or cache, shelve, dill, cloudpickle on data another party writes.
• Prefer safetensors, ONNX, skops, JSON, Parquet.

A finding when the file can come from outside the trust boundary: an upload, a URL, a user-chosen hub id, a shared bucket other tenants write. A model the team trains and ships in its own artifact store is not attacker input.

2. Code Evaluation

• eval, exec, compile, __import__/importlib.import_module with input.
• DataFrame.query(expr) / pd.eval(expr) with a user-built expression — evaluation, not filtering.
• Format-string injection: a user-supplied format string can walk attributes.

# ❌ template = "{user.__class__.__init__.__globals__[SECRET_KEY]}"
message = template.format(user=user)
# ✅ user text is the value, never the format
message = "Hello {name}".format(name=user.name)