Verify ML model integrity with checksums before loading to prevent tampering

secure-model-loader.py

"""Verify model integrity before loading.

Source: https://williamzujkowski.github.io/posts/2025-04-10-securing-personal-ai-experiments/

Two things this file exists to get right, because both are easy to get wrong in
a way that looks correct:

1. FAIL CLOSED. A hash check that skips verification for models it has never
   seen will happily load anything an attacker drops in the directory. If the
   hash file is missing entirely, it will load everything. A checker that
   passes what it does not recognise is worse than no checker, because you
   stop looking at that step.

2. torch.load UNPICKLES, AND UNPICKLING IS ARBITRARY CODE EXECUTION.
   PyTorch flipped the default to weights_only=True in 2.6 for exactly this
   reason. A checkpoint downloaded from a stranger is a program. Prefer
   safetensors; where a torch checkpoint is unavoidable, pass weights_only
   explicitly rather than trusting the installed version's default.
"""

import hashlib
import hmac
from pathlib import Path
from typing import Any, Dict


class SecureModelLoader:
    def __init__(self, trusted_hashes_file: str = "model_hashes.txt",
                 model_dir: str = "./models"):
        self.hashes_path = Path(trusted_hashes_file)
        self.model_dir = Path(model_dir).resolve()
        self.trusted_hashes = self._load_trusted_hashes()

    def _load_trusted_hashes(self) -> Dict[str, str]:
        # A missing hash file is a configuration error, not an empty allowlist.
        if not self.hashes_path.exists():
            raise FileNotFoundError(
                f"No trusted-hash file at {self.hashes_path}. Refusing to run "
                f"with an empty allowlist — every model would verify."
            )
        hashes = {}
        with open(self.hashes_path) as f:
            for line in f:
                line = line.strip()
                if not line or line.startswith("#"):
                    continue
                name, checksum = line.split(":", 1)
                hashes[name.strip()] = checksum.strip()
        if not hashes:
            raise ValueError(f"{self.hashes_path} contains no entries.")
        return hashes

    def resolve_model_path(self, name: str) -> Path:
        """Resolve inside model_dir and prove containment.

        Path traversal is prevented by resolving and checking containment, not
        by string-replacing '../' — that is trivially defeated, since '....//'
        collapses back to '../' after one non-recursive pass.
        """
        candidate = (self.model_dir / name).resolve()
        if not candidate.is_relative_to(self.model_dir):
            raise ValueError(f"Path escapes model directory: {name}")
        return candidate

    def verify_model(self, model_path: Path) -> None:
        """Raise unless the file matches a recorded hash. Fails closed."""
        expected = self.trusted_hashes.get(model_path.name)
        if expected is None:
            raise ValueError(
                f"No trusted hash on record for {model_path.name}. "
                f"Add one deliberately after verifying provenance."
            )

        digest = hashlib.sha256()
        with open(model_path, "rb") as f:
            for block in iter(lambda: f.read(1024 * 1024), b""):
                digest.update(block)

        if not hmac.compare_digest(digest.hexdigest(), expected):
            raise ValueError(f"Checksum mismatch for {model_path.name}")

    def load_safetensors(self, name: str) -> Dict[str, Any]:
        """Preferred path. safetensors cannot execute code by construction."""
        from safetensors.torch import load_file

        path = self.resolve_model_path(name)
        self.verify_model(path)
        return load_file(str(path))

    def load_torch_checkpoint(self, name: str) -> Any:
        """Fallback for .pt/.pth. weights_only is passed explicitly."""
        import torch

        path = self.resolve_model_path(name)
        self.verify_model(path)
        return torch.load(path, map_location="cpu", weights_only=True)