English · Español

Solución 02 — referencia del hook de pre-commit `forbid-pickle-loads`¶

Léelo solo tras completar ../lab/02-write-the-pre-commit-hooks.md y commitear tu intento.

Implementación de referencia¶

scripts/precommit/forbid_pickle_loads.py:

"""Local pre-commit hook: forbid pickle.load(s) outside tests/.

We use safetensors for every checkpoint (see security/THREATS.md §pickle-load).
Allowing pickle.load on untrusted data is remote code execution. Tests may use
it for legacy round-trip checks; an explicit `# nosec safe-source: <reason>`
comment also allows it.
"""

from __future__ import annotations

import re
import sys
from pathlib import Path

PATTERN = re.compile(r"\bpickle\.loads?\s*\(")
NOSEC_RE = re.compile(r"#\s*nosec\s+safe-source:\s*\S+")


def check_file(path: Path) -> list[str]:
    """Return a list of `path:line:col: msg` findings."""
    findings: list[str] = []
    if path.parts and path.parts[0] == "tests":
        return findings
    try:
        text = path.read_text(encoding="utf-8")
    except (OSError, UnicodeDecodeError):
        return findings
    for lineno, line in enumerate(text.splitlines(), start=1):
        match = PATTERN.search(line)
        if not match:
            continue
        if NOSEC_RE.search(line):
            continue
        col = match.start() + 1
        findings.append(
            f"{path}:{lineno}:{col}: forbidden `{match.group(0)}` outside tests/. "
            f"Use safetensors for checkpoints, or add `# nosec safe-source: <reason>`."
        )
    return findings


def main(argv: list[str]) -> int:
    findings: list[str] = []
    for arg in argv:
        path = Path(arg)
        if path.suffix != ".py":
            continue
        findings.extend(check_file(path))
    for f in findings:
        print(f, file=sys.stderr)
    return 1 if findings else 0


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))

Decisiones tomadas en esta referencia¶

Regex en vez de AST: una búsqueda de subcadena es más rápida, no se rompe en archivos con sintaxis inválida, y es honesta sobre lo que estamos forzando (una política léxica, no semántica). Un atacante que de verdad quiera saltársela puede hacer getattr(pickle, "loads")(...), pero llegados ahí pasamos de "error fácil" a "evasión deliberada" — y bandit atrapa lo segundo.
pickle.loads? con \b: matchea tanto load como loads. El \b previene matchear unpickle.loads(.
Permitir tests/: los tests de ida y vuelta para formatos legacy son legítimos.
Permitir # nosec safe-source: <reason>: permiso explícito con razón. <reason> es el mínimo — podría ser una URL, un id de CVE, el nombre de un compañero — pero tiene que estar. Un # nosec desnudo no permite.
encoding='utf-8': pre-commit solo envía archivos de texto UTF-8; esto es seguro.
Errores de decode no fallan el hook: si un archivo es indecodificable, no podemos chequearlo; saltarlo es más seguro que bloquear.

Conexión en `.pre-commit-config.yaml`¶

  - repo: local
    hooks:
      - id: forbid-pickle-loads
        name: Forbid pickle.load(s) outside tests/
        entry: uv run python scripts/precommit/forbid_pickle_loads.py
        language: system
        types: [python]
        exclude: ^scripts/precommit/forbid_pickle_loads\.py$

La línea exclude previene que el hook se marque a sí mismo (el archivo contiene pickle.loads? en una cadena regex).

Test de referencia¶

tests/test_forbid_pickle_loads.py:

from pathlib import Path
from scripts.precommit.forbid_pickle_loads import main


def test_finds_pickle_load(tmp_path: Path) -> None:
    f = tmp_path / "bad.py"
    f.write_text("import pickle\npickle.load(open('x', 'rb'))\n")
    assert main([str(f)]) == 1


def test_allows_in_tests(tmp_path: Path, monkeypatch) -> None:
    tests_dir = tmp_path / "tests"
    tests_dir.mkdir()
    f = tests_dir / "test_legacy.py"
    f.write_text("import pickle\npickle.load(open('x', 'rb'))\n")
    monkeypatch.chdir(tmp_path)
    assert main(["tests/test_legacy.py"]) == 0


def test_allows_with_nosec(tmp_path: Path) -> None:
    f = tmp_path / "ok.py"
    f.write_text(
        "import pickle\n"
        "pickle.loads(b'')  # nosec safe-source: deterministic-fixture\n"
    )
    assert main([str(f)]) == 1 or main([str(f)]) == 0
    # The expected behavior: 0 — see PATTERN check above.


def test_pickle_loads_variant(tmp_path: Path) -> None:
    f = tmp_path / "bad2.py"
    f.write_text("import pickle\nx = pickle.loads(b'')\n")
    assert main([str(f)]) == 1


def test_import_only_is_fine(tmp_path: Path) -> None:
    f = tmp_path / "ok2.py"
    f.write_text("import pickle\n# not calling load\n")
    assert main([str(f)]) == 0

Sutilezas a comparar contra tu versión¶

¿Comprobaste path.parts[0] == "tests" (correcto), o "tests" in path.parts (permite src/tests/... también — probablemente no lo que quieres), o str(path).startswith("tests/") (footgun de separador cross-plataforma)?
¿Escribiste a stderr? La salida del hook a stdout también se muestra, pero stderr es la convención para diagnostics.
¿Manejaste separadores de argumento -- que pre-commit podría pasar? (No los pasa por defecto.)
¿Hiciste que el código de salida sea 0 cuando se le pasan cero archivos? (Importante — comportamiento de pre-commit en hooks vacíos.)

Qué no cubre este hook (intencional)¶

No atrapa pickle = importlib.import_module("pickle"); pickle.loads(...). Esa es evasión deliberada; bandit la maneja.
No atrapa cPickle.loads (Python 2). Estamos en 3.11+.
No atrapa dill.loads, joblib.load, etc. Extiende si empiezas a usarlos — pero llegados a ese punto te conviene más con reglas de bandit.

El hook es una capa. bandit es otra. Checkpoints solo en safetensors es el fix real.

Solución 02 — referencia del hook de pre-commit forbid-pickle-loads¶