Fix model downloads on machines with a broken system trust store

setup.py

@@ -51,6 +51,7 @@ def get_string(string, rel_path="src/stepcount/__init__.py"):
     include_package_data=False,
     install_requires=[
         "actipy>=3.8.0",
+        "certifi>=2024.7.4",  # CA bundle for the HTTPS fallback; floor clears CVE-2024-39689
         "numpy==1.24.*",
         "scipy==1.10.*",
         "pandas==2.0.*",

src/stepcount/stepcount.py

@@ -4,6 +4,7 @@
 import pathlib
 import urllib
 import shutil
+import ssl
 import time
 import argparse
 import json
@@ -99,6 +100,9 @@ def main():
     parser.add_argument('--quiet', '-q', action='store_true', help='Suppress output')
     args = parser.parse_args()
 
+    # Arm the certificate fallback before any model download can run.
+    _ensure_download_ssl_context(verbose=not args.quiet)
+
     if args.download_models:
         download_models(force_download=args.force_download, ssl_repo_path=args.ssl_repo_path)
         return
@@ -477,6 +481,69 @@ def main():
     print(f"Done! ({round(after - before,2)}s)")
 
 
+def _ensure_download_ssl_context(verbose=True):
+    """Fall back to certifi's CA bundle when the default HTTPS context is unusable.
+
+    On affected OpenSSL builds (CVE-2026-34180), constructing the default SSL
+    context raises while loading the OS certificate store, which would otherwise
+    break every HTTPS download here — the model files and the torch.hub fetch.
+    certifi's bundle sidesteps the OS trust store. This engages only when the
+    default context is already broken and no custom hook is installed, so healthy
+    machines, and applications that configured their own trust, are untouched.
+    """
+    # A non-default hook means trust was configured deliberately elsewhere;
+    # respect it rather than overwriting it.
+    if ssl._create_default_https_context is not ssl.create_default_context:
+        return
+
+    try:
+        ssl.create_default_context()
+        return  # default context works; nothing to do
+    except ssl.SSLError:
+        pass
+
+    try:
+        import certifi
+    except ImportError:
+        return  # nothing we can do; let the original download error surface
+
+    def _certifi_https_context(*_args, **_kwargs):
+        return ssl.create_default_context(cafile=certifi.where())
+
+    # urllib and torch.hub both build their default context via this hook.
+    ssl._create_default_https_context = _certifi_https_context
+
+    if verbose:
+        print(
+            "Note: could not load the system certificate store "
+            "(known OpenSSL issue); using certifi CA bundle for downloads."
+        )
+
+
+def _download_to_file(url, dest, expected_md5=None, timeout=60):
+    """Download ``url`` to ``dest`` atomically.
+
+    Writes to a temp file and swaps it into place only once complete and (when
+    requested) MD5-verified, so an interrupted or corrupt download never leaves a
+    truncated file behind. The temp name is per-process so concurrent downloads
+    of the same file don't clobber each other, and a connection timeout keeps a
+    stalled server from hanging indefinitely.
+    """
+    dest = pathlib.Path(dest)
+    tmp_pth = dest.with_name(f"{dest.name}.{os.getpid()}.tmp")
+    try:
+        with urllib.request.urlopen(url, timeout=timeout) as f_src, open(tmp_pth, "wb") as f_dst:
+            shutil.copyfileobj(f_src, f_dst)
+        if expected_md5 is not None and utils.md5(tmp_pth) != expected_md5:
+            raise ValueError(
+                f"MD5 mismatch for downloaded file {dest.name}. Download may be corrupted."
+            )
+        os.replace(tmp_pth, dest)
+    except BaseException:
+        tmp_pth.unlink(missing_ok=True)
+        raise
+
+
 def download_models(force_download=False, ssl_repo_path=None):
     """Download all model files for offline use."""
 
@@ -485,19 +552,7 @@ def download_models(force_download=False, ssl_repo_path=None):
         if force_download or not pth.exists():
             url = f"https://wearables-files.ndph.ox.ac.uk/files/models/stepcount/{__model_version__[model_type]}.joblib.lzma"
             print(f"Downloading {url}...")
-            tmp_pth = pth.with_suffix('.tmp')
-            try:
-                with urllib.request.urlopen(url, timeout=60) as f_src, open(tmp_pth, "wb") as f_dst:
-                    shutil.copyfileobj(f_src, f_dst)
-                if utils.md5(tmp_pth) != __model_md5__[model_type]:
-                    raise ValueError(
-                        f"MD5 mismatch for {model_type} model. Download may be corrupted."
-                    )
-                os.replace(tmp_pth, pth)
-            except Exception:
-                if tmp_pth.exists():
-                    tmp_pth.unlink()
-                raise
+            _download_to_file(url, pth, expected_md5=__model_md5__[model_type])
             print(f"Saved to {pth}")
         else:
             print(f"Already exists: {pth}")
@@ -555,8 +610,10 @@ def load_model(
 
         print(f"Downloading {url}...")
 
-        with urllib.request.urlopen(url) as f_src, open(pth, "wb") as f_dst:
-            shutil.copyfileobj(f_src, f_dst)
+        _download_to_file(
+            url, pth,
+            expected_md5=__model_md5__[model_type] if check_md5 else None,
+        )
 
     if check_md5:
         assert utils.md5(pth) == __model_md5__[model_type], (

tests/test_stepcount.py

@@ -450,6 +450,172 @@ def test_load_model_missing_file(self, temp_dir):
         )
 
 
+class TestDownloadToFile:
+    """Tests for the atomic download helper `_download_to_file`."""
+
+    def test_success_writes_dest_and_cleans_temp_with_timeout(self, tmp_path, monkeypatch):
+        import io
+        import hashlib
+        import urllib.request
+        payload = b"model-payload-bytes"
+        dest = tmp_path / "model.joblib.lzma"
+        seen = {}
+
+        def fake_urlopen(url, timeout=None):
+            seen['timeout'] = timeout
+            return io.BytesIO(payload)
+
+        monkeypatch.setattr(urllib.request, "urlopen", fake_urlopen)
+
+        stepcount._download_to_file(
+            "http://example/model", dest,
+            expected_md5=hashlib.md5(payload).hexdigest(),
+        )
+
+        assert dest.read_bytes() == payload
+        assert seen['timeout'] == 60                   # a stalled server can't hang the download
+        assert list(tmp_path.glob("*.tmp")) == []      # no temp left behind
+
+    def test_md5_mismatch_raises_and_leaves_no_files(self, tmp_path, monkeypatch):
+        import io
+        import urllib.request
+        dest = tmp_path / "model.joblib.lzma"
+
+        monkeypatch.setattr(urllib.request, "urlopen",
+                            lambda url, timeout=None: io.BytesIO(b"corrupt"))
+
+        with pytest.raises(ValueError, match="MD5 mismatch"):
+            stepcount._download_to_file(
+                "http://example/model", dest, expected_md5="0" * 32)
+
+        assert not dest.exists()                        # bad download never lands at dest
+        assert list(tmp_path.glob("*.tmp")) == []       # temp cleaned up on failure
+
+    def test_failure_midstream_preserves_existing_dest(self, tmp_path, monkeypatch):
+        import os
+        import urllib.request
+        dest = tmp_path / "model.joblib.lzma"
+        dest.write_bytes(b"previous-good-model")        # a valid model already in place
+
+        class _BoomReader:
+            def __enter__(self):
+                return self
+
+            def __exit__(self, *a):
+                return False
+
+            def read(self, *a):
+                raise OSError("connection reset mid-download")
+
+        monkeypatch.setattr(urllib.request, "urlopen",
+                            lambda url, timeout=None: _BoomReader())
+
+        with pytest.raises(OSError):
+            stepcount._download_to_file(
+                "http://example/model", dest, expected_md5="0" * 32)
+
+        assert dest.read_bytes() == b"previous-good-model"  # a failed download can't corrupt a good file
+        # the per-process temp is cleaned up, nothing left behind
+        assert not (tmp_path / f"{dest.name}.{os.getpid()}.tmp").exists()
+        assert list(tmp_path.glob("*.tmp")) == []
+
+
+class TestEnsureDownloadSSLContext:
+    """Tests for the certifi SSL-context fallback used for model downloads."""
+
+    def test_noop_when_default_context_works(self, monkeypatch):
+        """When the active context factory builds fine, the HTTPS hook is untouched."""
+        import ssl
+
+        def ok_factory(*a, **k):
+            return "ok-context-sentinel"
+
+        # Force the "works" branch host-independently: make the default factory
+        # succeed, and keep the active hook identical to it so the
+        # default-in-effect guard holds. (Without this the test is coupled to the
+        # host's real trust store and would spuriously fail on a broken one.)
+        monkeypatch.setattr(ssl, 'create_default_context', ok_factory)
+        monkeypatch.setattr(ssl, '_create_default_https_context', ok_factory)
+
+        stepcount._ensure_download_ssl_context(verbose=False)
+
+        assert ssl._create_default_https_context is ok_factory
+
+    def test_falls_back_to_certifi_when_store_broken(self, monkeypatch):
+        """A broken default store installs a certifi-backed hook that still verifies."""
+        import ssl
+        certifi = pytest.importorskip('certifi')
+
+        real_create = ssl.create_default_context
+        seen_cafiles = []
+
+        def spy(*args, **kwargs):
+            seen_cafiles.append(kwargs.get('cafile'))
+            # Simulate the broken Windows store: the no-arg default build fails,
+            # but building from an explicit cafile (what the fallback does) works.
+            if not kwargs.get('cafile'):
+                raise ssl.SSLError("[ASN1: NOT_ENOUGH_DATA] not enough data")
+            return real_create(*args, **kwargs)
+
+        # Model the stdlib default hook being active (identity holds) but broken.
+        monkeypatch.setattr(ssl, 'create_default_context', spy)
+        monkeypatch.setattr(ssl, '_create_default_https_context', spy)
+
+        stepcount._ensure_download_ssl_context(verbose=False)
+
+        hook = ssl._create_default_https_context
+        assert hook is not spy                       # a new fallback hook was installed
+        ctx = hook()                                 # exercise it, not just its identity
+        assert isinstance(ctx, ssl.SSLContext)
+        assert certifi.where() in seen_cafiles       # fallback uses certifi's CA bundle
+        assert ctx.check_hostname is True            # verification is preserved...
+        assert ctx.verify_mode == ssl.CERT_REQUIRED  # ...not silently downgraded
+
+    def test_preserves_custom_hook_when_store_broken(self, monkeypatch):
+        """A custom HTTPS hook installed by an embedding app is not overwritten."""
+        import ssl
+
+        def _boom(*a, **k):
+            raise ssl.SSLError("[ASN1: NOT_ENOUGH_DATA] not enough data")
+
+        def custom_hook(*a, **k):
+            return "custom-context-sentinel"
+
+        # Default factory is broken, but a distinct custom hook is already active.
+        monkeypatch.setattr(ssl, 'create_default_context', _boom)
+        monkeypatch.setattr(ssl, '_create_default_https_context', custom_hook)
+
+        stepcount._ensure_download_ssl_context(verbose=False)
+
+        # The custom hook must be respected, never silently replaced by certifi.
+        assert ssl._create_default_https_context is custom_hook
+
+    def test_no_fallback_without_certifi(self, monkeypatch):
+        """If certifi is unavailable, leave the HTTPS hook alone (surface original error)."""
+        import ssl
+        import builtins
+
+        def _boom(*a, **k):
+            raise ssl.SSLError("[ASN1: NOT_ENOUGH_DATA] not enough data")
+
+        # Model the default hook being active (identity holds) but the store broken.
+        monkeypatch.setattr(ssl, 'create_default_context', _boom)
+        monkeypatch.setattr(ssl, '_create_default_https_context', _boom)
+
+        real_import = builtins.__import__
+
+        def _no_certifi(name, *args, **kwargs):
+            if name == 'certifi':
+                raise ImportError("No module named 'certifi'")
+            return real_import(name, *args, **kwargs)
+
+        monkeypatch.setattr(builtins, '__import__', _no_certifi)
+        stepcount._ensure_download_ssl_context(verbose=False)
+
+        # certifi import failed, so no fallback was installed — hook is unchanged.
+        assert ssl._create_default_https_context is _boom
+
+
 class TestENMOCalculation:
     """Tests for ENMO (Euclidean Norm Minus One) calculation."""
 
@@ -584,7 +750,10 @@ def test_cli_download_models_in_help(self):
     def test_cli_download_models_no_filepath(self):
         """Test that --download-models works without a filepath argument."""
         from unittest.mock import patch
-        with patch('stepcount.stepcount.download_models') as mock_dl:
+        # Stub the SSL setup too, so main() can't mutate global ssl state on a
+        # broken-store host and leak the process-global hook into later tests.
+        with patch('stepcount.stepcount.download_models') as mock_dl, \
+                patch('stepcount.stepcount._ensure_download_ssl_context'):
             # Simulate calling main() with --download-models
             with patch('sys.argv', ['stepcount', '--download-models']):
                 stepcount.main()