test(security): cover shell capabilities and updater trust

backend/tests/test_security_hardening_2026_08.py

@@ -0,0 +1,130 @@
+from __future__ import annotations
+
+import asyncio
+import hashlib
+import inspect
+from types import SimpleNamespace
+
+import pytest
+from fastapi import WebSocketDisconnect
+
+from routers import agent_shell
+from services import agent_shell_ws_token as ws_tokens
+from services import updater
+
+
+class _FakeWebSocket:
+    def __init__(self, headers: dict[str, str] | None = None, host: str = "127.0.0.1") -> None:
+        self.headers = headers or {}
+        self.client = SimpleNamespace(host=host)
+        self.closed: tuple[int, str] | None = None
+
+    async def close(self, code: int = 1000, reason: str = "") -> None:
+        self.closed = (code, reason)
+
+
+def setup_function() -> None:
+    ws_tokens.reset_agent_shell_ws_tokens_for_tests()
+
+
+def teardown_function() -> None:
+    ws_tokens.reset_agent_shell_ws_tokens_for_tests()
+
+
+def test_agent_shell_query_no_longer_accepts_admin_key() -> None:
+    parameters = inspect.signature(agent_shell.agent_shell_websocket).parameters
+    assert "admin_key" not in parameters
+    assert "ws_token" in parameters
+
+
+def test_loopback_source_ip_is_not_agent_shell_authorization() -> None:
+    ws = _FakeWebSocket({"host": "127.0.0.1:8000"})
+    with pytest.raises(WebSocketDisconnect):
+        asyncio.run(agent_shell._authorize_agent_shell_ws(ws, ""))
+    assert ws.closed is not None
+    assert ws.closed[0] == 4403
+
+
+def test_browser_one_time_token_allows_same_host_origin() -> None:
+    token, _ = ws_tokens.mint_agent_shell_ws_token()
+    ws = _FakeWebSocket(
+        {
+            "host": "127.0.0.1:8000",
+            "origin": "http://127.0.0.1:3000",
+        }
+    )
+    asyncio.run(agent_shell._authorize_agent_shell_ws(ws, token))
+
+    # The capability is single-use.
+    second = _FakeWebSocket(
+        {
+            "host": "127.0.0.1:8000",
+            "origin": "http://127.0.0.1:3000",
+        }
+    )
+    with pytest.raises(WebSocketDisconnect):
+        asyncio.run(agent_shell._authorize_agent_shell_ws(second, token))
+
+
+def test_cross_origin_browser_is_rejected_before_token_use() -> None:
+    token, _ = ws_tokens.mint_agent_shell_ws_token()
+    ws = _FakeWebSocket(
+        {
+            "host": "127.0.0.1:8000",
+            "origin": "https://attacker.example",
+        }
+    )
+    with pytest.raises(WebSocketDisconnect):
+        asyncio.run(agent_shell._authorize_agent_shell_ws(ws, token))
+    assert ws.closed is not None
+    assert ws.closed[0] == 4403
+    # Rejection happens before capability consumption.
+    assert ws_tokens.consume_agent_shell_ws_token(token) is True
+
+
+def test_agent_shell_token_store_stays_bounded_without_recursive_locking() -> None:
+    minted = [ws_tokens.mint_agent_shell_ws_token()[0] for _ in range(400)]
+    assert len(ws_tokens._store) <= ws_tokens._MAX_ACTIVE_TOKENS
+    assert ws_tokens.consume_agent_shell_ws_token(minted[-1]) is True
+
+
+def test_updater_refuses_same_release_checksum_as_sole_trust_root(tmp_path, monkeypatch) -> None:
+    archive = tmp_path / "ShadowBroker_v9.9.9.zip"
+    archive.write_bytes(b"archive-under-test")
+    digest = hashlib.sha256(archive.read_bytes()).hexdigest()
+
+    monkeypatch.delenv("MESH_UPDATE_SHA256", raising=False)
+    monkeypatch.setattr(updater, "_load_baked_in_release_digests", lambda: {})
+    monkeypatch.setattr(
+        updater,
+        "_fetch_sha256sums",
+        lambda _url: {archive.name: digest},
+    )
+
+    with pytest.raises(RuntimeError, match="no independent archive digest"):
+        updater._validate_zip_hash(
+            str(archive),
+            asset_name=archive.name,
+            sha256sums_url="https://github.com/BigBodyCobain/Shadowbroker/releases/download/v9.9.9/SHA256SUMS.txt",
+            release_tag="v9.9.9",
+        )
+
+
+def test_updater_accepts_preinstalled_baked_digest(tmp_path, monkeypatch) -> None:
+    archive = tmp_path / "ShadowBroker_v9.9.9.zip"
+    archive.write_bytes(b"trusted-archive")
+    digest = hashlib.sha256(archive.read_bytes()).hexdigest()
+
+    monkeypatch.delenv("MESH_UPDATE_SHA256", raising=False)
+    monkeypatch.setattr(
+        updater,
+        "_load_baked_in_release_digests",
+        lambda: {"v9.9.9": {archive.name: digest}},
+    )
+
+    note = updater._validate_zip_hash(
+        str(archive),
+        asset_name=archive.name,
+        release_tag="v9.9.9",
+    )
+    assert "baked-in digest" in note