R-5: Verify SHA-256 checksums on downloaded model files

autoptz/engine/runtime/models.py

@@ -31,10 +31,12 @@
 
 from __future__ import annotations
 
+import hashlib
 import logging
 import os
 import tempfile
 import time
+import urllib.error
 import urllib.request
 from collections.abc import Callable
 from pathlib import Path
@@ -136,6 +138,16 @@ def detector_model_for_tier(tier: str | None) -> str:
 # error page / truncated file as a "model".
 _MIN_ONNX_BYTES = 1 << 18  # 256 KiB
 
+# Chunk size (bytes) used to hash a downloaded model without loading it whole
+# into memory.  Mirrors autoptz.update.installer.verify_sha256's approach —
+# copied rather than imported so the engine never depends on the updater
+# package (see module docstring / layering).
+_HASH_CHUNK_SIZE = 1024 * 1024
+
+# Name of the checksum manifest published next to prebuilt model assets, in
+# the same release/directory as the ``.onnx`` files themselves.
+_CHECKSUM_MANIFEST_NAME = "SHA256SUMS"
+
 # Export knobs that match what detect.py expects (see module docstring).
 _EXPORT_KWARGS = {"format": "onnx", "nms": False, "dynamic": False, "opset": 12}
 _DISABLE_EXPORT_ENV = "AUTOPTZ_NO_MODEL_EXPORT"
@@ -235,6 +247,60 @@ def bundled_models_dir() -> Path:
     return Path(__file__).resolve().parents[2] / "models"
 
 
+def _sha256_of(path: Path) -> str:
+    """Return the lower-case hex SHA-256 digest of *path*, read in chunks.
+
+    Mirrors ``autoptz.update.installer.verify_sha256``'s chunked-hashing
+    approach (see module docstring); duplicated locally so the engine never
+    imports the updater package.
+    """
+    h = hashlib.sha256()
+    with path.open("rb") as f:
+        while True:
+            chunk = f.read(_HASH_CHUNK_SIZE)
+            if not chunk:
+                break
+            h.update(chunk)
+    return h.hexdigest().lower()
+
+
+def _manifest_url_for(model_url: str) -> str:
+    """Return the ``SHA256SUMS`` manifest URL published next to *model_url*.
+
+    The manifest is expected in the same "directory" as the model asset —
+    i.e. the same base URL with the filename replaced.
+    """
+    base, _, _name = model_url.rpartition("/")
+    if not base:
+        return ""
+    return f"{base}/{_CHECKSUM_MANIFEST_NAME}"
+
+
+def _parse_sha256sums_manifest(content: str, filename: str) -> str | None:
+    """Extract the hex digest for *filename* from ``SHA256SUMS``-style *content*.
+
+    Expects the standard ``<hex>  <filename>`` (or single-space) format used by
+    ``sha256sum``.  Parses defensively: blank lines, extra whitespace, and a
+    leading ``*`` (binary-mode marker) are tolerated.  Returns ``None`` if the
+    file has no entry for *filename*.
+    """
+    target = filename.strip().lower()
+    for line in content.splitlines():
+        line = line.strip()
+        if not line or line.startswith("#"):
+            continue
+        parts = line.split(None, 1)
+        if len(parts) != 2:
+            continue
+        digest, name = parts
+        name = name.strip().lstrip("*").strip().lower()
+        # Tolerate a manifest that lists paths ("models/yolo11n.onnx").
+        basename = name.replace("\\", "/").rsplit("/", 1)[-1]
+        if name == target or basename == target:
+            return digest.strip().lower()
+    return None
+
+
 class ModelManager:
     """Resolves / downloads / exports the models the engine needs.
 
@@ -671,6 +737,9 @@ def _download_prebuilt(self, onnx_path: Path) -> str | None:
                 )
                 return None
 
+            if not self._verify_checksum(tmp_path, url, onnx_path.name):
+                return None
+
             _replace_with_retry(tmp_path, onnx_path)
             log.info(
                 "Model ONNX ready at %s (%.1f MB, prebuilt, torch-free)",
@@ -691,6 +760,74 @@ def _download_prebuilt(self, onnx_path: Path) -> str | None:
             except Exception:  # noqa: BLE001
                 log.debug("Could not remove temp download %s", tmp_path, exc_info=True)
 
+    def _verify_checksum(self, tmp_path: Path, model_url: str, asset_name: str) -> bool:
+        """Verify *tmp_path* against the release's ``SHA256SUMS`` manifest.
+
+        Fetched once per download from the same base URL as the model asset
+        itself.  Three outcomes:
+
+        - Manifest unreachable (404/network error) or parses to no entry for
+          *asset_name*: log a ``WARNING`` ("legacy release without
+          checksums") and return ``True`` (proceed) — hard-failing on a
+          missing manifest is deliberately deferred (see module docstring).
+        - Entry present and the digest matches: log ``DEBUG``/``INFO`` and
+          return ``True``.
+        - Entry present and the digest does NOT match: delete *tmp_path*, log
+          an ``ERROR``, and return ``False`` so the caller treats this as a
+          failed download (existing fallback machinery applies — no new
+          machinery is introduced here).
+        """
+        manifest_url = _manifest_url_for(model_url)
+        if not manifest_url:
+            log.warning(
+                "Could not derive a SHA256SUMS manifest URL from %s; "
+                "skipping integrity check for %s (legacy release without checksums).",
+                model_url,
+                asset_name,
+            )
+            return True
+
+        try:
+            with urllib.request.urlopen(manifest_url) as resp:  # noqa: S310
+                manifest_text = resp.read().decode("utf-8", errors="replace")
+        except (OSError, urllib.error.URLError, TimeoutError) as exc:
+            log.warning(
+                "No SHA256SUMS manifest found at %s (%s) — proceeding without "
+                "integrity verification for %s (legacy release without checksums).",
+                manifest_url,
+                exc,
+                asset_name,
+            )
+            return True
+
+        expected = _parse_sha256sums_manifest(manifest_text, asset_name)
+        if expected is None:
+            log.warning(
+                "SHA256SUMS manifest at %s has no entry for %s — proceeding without "
+                "integrity verification (legacy release without checksums).",
+                manifest_url,
+                asset_name,
+            )
+            return True
+
+        actual = _sha256_of(tmp_path)
+        if actual != expected:
+            try:
+                tmp_path.unlink()
+            except OSError:
+                log.debug("Could not remove mismatched download %s", tmp_path, exc_info=True)
+            log.error(
+                "SHA-256 mismatch for %s: expected %s, got %s — download deleted, "
+                "not using this file.",
+                asset_name,
+                expected,
+                actual,
+            )
+            return False
+
+        log.info("SHA-256 verified for %s", asset_name)
+        return True
+
     def _download_and_export(self, model_pt: str, onnx_path: Path) -> str | None:
         """Download the ultralytics ``.pt`` and export it to *onnx_path*.
 

tests/test_discovery.py

@@ -535,8 +535,11 @@ def test_stable_device_not_re_reported(self) -> None:
         added = [e for e in events if e[0] == "added"]
         assert len(added) == 1
 
-    def test_wsdiscovery_unavailable_does_not_raise(self) -> None:
-        _remove_mock_wsdiscovery()
+    def test_wsdiscovery_unavailable_does_not_raise(self, monkeypatch) -> None:
+        # None in sys.modules makes `import wsdiscovery` raise ImportError even
+        # when the real package is installed — otherwise this test runs a live
+        # WS-Discovery scan and fails on any LAN with responsive devices.
+        monkeypatch.setitem(sys.modules, "wsdiscovery", None)
 
         discovery = ONVIFDiscovery(rescan_interval=0.05)
         discovery.start()

tests/test_models.py

@@ -9,8 +9,10 @@
 
 from __future__ import annotations
 
+import hashlib
 import sys
 import types
+import urllib.error
 from pathlib import Path
 from unittest.mock import MagicMock, patch
 
@@ -412,6 +414,9 @@ def test_prebuilt_download_is_preferred_over_export(tmp_path, monkeypatch) -> No
     payload = b"\x00" * (300 * 1024)  # > _MIN_ONNX_BYTES (256 KiB)
 
     def fake_urlopen(url, *a, **k):
+        if url.endswith("/SHA256SUMS"):
+            # No manifest published for this release — legacy path, proceeds.
+            raise urllib.error.HTTPError(url, 404, "Not Found", {}, None)
         assert url == "https://example.test/yolo11n.onnx"
         return _FakeHTTPResponse(payload)
 
@@ -536,6 +541,179 @@ def fake_urlopen(url, *a, **k):
 # ── cache dir resolution ──────────────────────────────────────────────────────
 
 
+# ── SHA-256 manifest verification (prebuilt downloads) ────────────────────────
+
+
+def _sha256_hex(data: bytes) -> str:
+    return hashlib.sha256(data).hexdigest()
+
+
+def _make_urlopen(*, model_url: str, model_payload: bytes, manifest: str | None):
+    """Return a fake ``urlopen`` serving *model_url* and a sibling ``SHA256SUMS``.
+
+    ``manifest=None`` simulates a 404 (no checksum manifest published for this
+    release yet); any other string is served verbatim for the ``SHA256SUMS``
+    URL derived from *model_url*'s directory.
+    """
+    manifest_url = model_url.rsplit("/", 1)[0] + "/SHA256SUMS"
+
+    def fake_urlopen(url, *a, **k):
+        if url == model_url:
+            return _FakeHTTPResponse(model_payload)
+        if url == manifest_url:
+            if manifest is None:
+                raise urllib.error.HTTPError(url, 404, "Not Found", {}, None)
+            return _FakeHTTPResponse(manifest.encode("utf-8"))
+        raise AssertionError(f"unexpected URL: {url}")
+
+    return fake_urlopen
+
+
+def test_prebuilt_download_good_checksum_is_kept(tmp_path, monkeypatch, caplog) -> None:
+    """(a) Manifest present + entry matches → file is verified and kept."""
+    monkeypatch.delenv("AUTOPTZ_MODEL_PATH", raising=False)
+    model_url = "https://example.test/yolo11n.onnx"
+    monkeypatch.setenv("AUTOPTZ_MODEL_URL", model_url)
+    monkeypatch.setitem(sys.modules, "ultralytics", None)
+
+    payload = b"\x00" * (300 * 1024)
+    manifest = f"{_sha256_hex(payload)}  yolo11n.onnx\n"
+    fake_urlopen = _make_urlopen(model_url=model_url, model_payload=payload, manifest=manifest)
+
+    cache = tmp_path / "cache"
+    with patch("autoptz.engine.runtime.models.urllib.request.urlopen", fake_urlopen):
+        with caplog.at_level("DEBUG"):
+            mgr = ModelManager(cache_dir=cache)
+            result = mgr.ensure_detector()
+
+    onnx = cache / "yolo11n.onnx"
+    assert result == str(onnx)
+    assert onnx.is_file()
+    assert onnx.read_bytes() == payload
+
+
+def test_prebuilt_download_bad_checksum_is_deleted_and_errors(
+    tmp_path, monkeypatch, caplog
+) -> None:
+    """(b) Manifest present + mismatch → file deleted, ERROR logged, no fallback machinery."""
+    monkeypatch.delenv("AUTOPTZ_MODEL_PATH", raising=False)
+    model_url = "https://example.test/yolo11n.onnx"
+    monkeypatch.setenv("AUTOPTZ_MODEL_URL", model_url)
+    # No ultralytics → if the code fell back to export, it would return None too,
+    # so also assert the tmp/final file never lingers to prove the delete path ran.
+    monkeypatch.setitem(sys.modules, "ultralytics", None)
+
+    payload = b"\x00" * (300 * 1024)
+    wrong_hash = "0" * 64
+    manifest = f"{wrong_hash}  yolo11n.onnx\n"
+    fake_urlopen = _make_urlopen(model_url=model_url, model_payload=payload, manifest=manifest)
+
+    cache = tmp_path / "cache"
+    with patch("autoptz.engine.runtime.models.urllib.request.urlopen", fake_urlopen):
+        with caplog.at_level("ERROR"):
+            mgr = ModelManager(cache_dir=cache)
+            result = mgr.ensure_detector()
+
+    onnx = cache / "yolo11n.onnx"
+    assert result is None
+    assert not onnx.exists()
+    assert not any(p.suffix == ".part" for p in cache.glob("*.part"))
+    assert any(
+        "SHA-256" in rec.message or "checksum" in rec.message.lower() for rec in caplog.records
+    )
+
+
+def test_prebuilt_download_missing_manifest_warns_and_keeps_file(
+    tmp_path, monkeypatch, caplog
+) -> None:
+    """(c) Manifest missing (404) → WARNING logged, file kept (legacy release)."""
+    monkeypatch.delenv("AUTOPTZ_MODEL_PATH", raising=False)
+    model_url = "https://example.test/yolo11n.onnx"
+    monkeypatch.setenv("AUTOPTZ_MODEL_URL", model_url)
+    monkeypatch.setitem(sys.modules, "ultralytics", None)
+
+    payload = b"\x00" * (300 * 1024)
+    fake_urlopen = _make_urlopen(model_url=model_url, model_payload=payload, manifest=None)
+
+    cache = tmp_path / "cache"
+    with patch("autoptz.engine.runtime.models.urllib.request.urlopen", fake_urlopen):
+        with caplog.at_level("WARNING"):
+            mgr = ModelManager(cache_dir=cache)
+            result = mgr.ensure_detector()
+
+    onnx = cache / "yolo11n.onnx"
+    assert result == str(onnx)
+    assert onnx.is_file()
+    assert onnx.read_bytes() == payload
+    assert any(
+        "legacy release" in rec.message.lower() or "no checksum" in rec.message.lower()
+        for rec in caplog.records
+        if rec.levelname == "WARNING"
+    )
+
+
+def test_prebuilt_download_entry_absent_warns_and_keeps_file(tmp_path, monkeypatch, caplog) -> None:
+    """(d) Manifest present but has no entry for this file → WARNING, file kept."""
+    monkeypatch.delenv("AUTOPTZ_MODEL_PATH", raising=False)
+    model_url = "https://example.test/yolo11n.onnx"
+    monkeypatch.setenv("AUTOPTZ_MODEL_URL", model_url)
+    monkeypatch.setitem(sys.modules, "ultralytics", None)
+
+    payload = b"\x00" * (300 * 1024)
+    # Manifest published, but only lists an unrelated file.
+    manifest = f"{_sha256_hex(b'other')}  yolo11s.onnx\n"
+    fake_urlopen = _make_urlopen(model_url=model_url, model_payload=payload, manifest=manifest)
+
+    cache = tmp_path / "cache"
+    with patch("autoptz.engine.runtime.models.urllib.request.urlopen", fake_urlopen):
+        with caplog.at_level("WARNING"):
+            mgr = ModelManager(cache_dir=cache)
+            result = mgr.ensure_detector()
+
+    onnx = cache / "yolo11n.onnx"
+    assert result == str(onnx)
+    assert onnx.is_file()
+    assert onnx.read_bytes() == payload
+    assert any(rec.levelname == "WARNING" for rec in caplog.records)
+
+
+def test_fetch_models_tool_path_verifies_checksum(tmp_path, monkeypatch) -> None:
+    """(e) tools/fetch_models.py's ModelManager.ensure_app_models exercises the
+    same checksum helper — a bad checksum must not leave a bad file cached even
+    when driven through the CLI's entry point (ensure_app_models -> ensure_detector
+    -> ensure_pose, all routed through ModelManager._download_prebuilt).
+    """
+    monkeypatch.delenv("AUTOPTZ_MODEL_PATH", raising=False)
+    monkeypatch.delenv("AUTOPTZ_POSE_MODEL_PATH", raising=False)
+    monkeypatch.setenv("AUTOPTZ_MODEL_URL", "https://example.test/{stem}.onnx")
+    monkeypatch.setitem(sys.modules, "ultralytics", None)
+
+    payload = b"\x00" * (300 * 1024)
+    good_hash = _sha256_hex(payload)
+
+    def fake_urlopen(url, *a, **k):
+        if url.endswith("/SHA256SUMS"):
+            # Only yolo11n (fast tier) gets a correct entry; everything else is
+            # either wrong or absent, to prove each is verified independently.
+            manifest = f"{good_hash}  yolo11n.onnx\n{'0' * 64}  yolo11s.onnx\n"
+            return _FakeHTTPResponse(manifest.encode("utf-8"))
+        return _FakeHTTPResponse(payload)
+
+    cache = tmp_path / "cache"
+    with patch("autoptz.engine.runtime.models.urllib.request.urlopen", fake_urlopen):
+        mgr = ModelManager(cache_dir=cache)
+        rows = mgr.ensure_app_models(
+            keys=["detector_fast", "detector_balanced"], include_pose=False
+        )
+
+    by_key = {r["key"]: r for r in rows}
+    assert by_key["detector_fast"]["state"] == "ok"
+    assert (cache / "yolo11n.onnx").is_file()
+    # yolo11s.onnx had a checksum mismatch and no ultralytics fallback → failed.
+    assert by_key["detector_balanced"]["state"] == "failed"
+    assert not (cache / "yolo11s.onnx").exists()
+
+
 def test_default_cache_dir_is_under_appdata_models() -> None:
     mgr = ModelManager()
     # Lives under the platform AutoPTZ dir, in a "models" subfolder.