diff --git a/apps/device-host-agent/host_agent/mcp_token.py b/apps/device-host-agent/host_agent/mcp_token.py new file mode 100644 index 0000000..a4cb0ca --- /dev/null +++ b/apps/device-host-agent/host_agent/mcp_token.py @@ -0,0 +1,118 @@ +"""Bearer-token persistence for the host-agent MCP server. + +The token is generated on first start and persisted to a JSON file with +0o600 permissions (POSIX) alongside the host identity. Rotation = delete +the file and restart host-agent. +""" + +from __future__ import annotations + +import json +import os +import secrets +import tempfile +from dataclasses import dataclass +from datetime import UTC, datetime +from pathlib import Path +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Callable + + +_TOKEN_BYTES = 32 + + +class McpTokenStoreError(RuntimeError): + """Raised when the MCP token file cannot be read or written.""" + + +@dataclass(frozen=True) +class McpToken: + version: int + token: str + created_at: datetime + + +class McpTokenStore: + def __init__( + self, + path: Path, + *, + now: Callable[[], datetime] | None = None, + ) -> None: + self._path = Path(path) + self._now = now or (lambda: datetime.now(UTC)) + + def load_or_create(self) -> McpToken: + if self._path.exists(): + return self._read_existing() + return self._generate_and_write() + + def verify(self, presented: str) -> bool: + try: + token = self.load_or_create() + except McpTokenStoreError: + return False + import hmac + + return hmac.compare_digest(token.token, presented) + + def _read_existing(self) -> McpToken: + try: + data = json.loads(self._path.read_text()) + except (OSError, json.JSONDecodeError) as exc: + raise McpTokenStoreError( + f"cannot read MCP token file {self._path}: {exc}" + ) from exc + if not isinstance(data, dict): + raise McpTokenStoreError("MCP token file is not a JSON object") + try: + return McpToken( + version=int(data["version"]), + token=str(data["token"]), + created_at=datetime.fromisoformat(str(data["created_at"])), + ) + except (KeyError, TypeError, ValueError) as exc: + raise McpTokenStoreError( + f"MCP token file schema invalid: {exc}" + ) from exc + + def _generate_and_write(self) -> McpToken: + token = McpToken( + version=1, + token=secrets.token_urlsafe(_TOKEN_BYTES), + created_at=self._now(), + ) + payload = { + "version": token.version, + "token": token.token, + "created_at": token.created_at.isoformat(), + } + try: + self._atomic_write(json.dumps(payload, indent=2)) + except OSError as exc: + raise McpTokenStoreError( + f"cannot write MCP token file {self._path}: {exc}" + ) from exc + return token + + def _atomic_write(self, content: str) -> None: + self._path.parent.mkdir(parents=True, exist_ok=True) + # Atomic on POSIX; on Windows os.replace is also atomic per docs. + fd, tmp_name = tempfile.mkstemp( + prefix=".host_mcp_token.", + suffix=".tmp", + dir=str(self._path.parent), + ) + try: + with os.fdopen(fd, "w", encoding="utf-8") as fh: + fh.write(content) + os.chmod(tmp_name, 0o600) + os.replace(tmp_name, self._path) + except BaseException: + try: + os.unlink(tmp_name) + except OSError: + pass + raise diff --git a/apps/device-host-agent/tests/test_mcp_token.py b/apps/device-host-agent/tests/test_mcp_token.py new file mode 100644 index 0000000..edcb617 --- /dev/null +++ b/apps/device-host-agent/tests/test_mcp_token.py @@ -0,0 +1,107 @@ +from __future__ import annotations + +import json +import os +import stat +import sys +from datetime import datetime +from pathlib import Path + +import pytest + +from host_agent.mcp_token import McpToken, McpTokenStore, McpTokenStoreError + + +def test_load_or_create_generates_when_missing(tmp_path: Path) -> None: + store = McpTokenStore(tmp_path / "host_mcp_token.json") + token = store.load_or_create() + assert token.version == 1 + assert len(token.token) >= 40 # secrets.token_urlsafe(32) -> ~43 chars + assert isinstance(token.created_at, datetime) + # File now exists. + assert (tmp_path / "host_mcp_token.json").exists() + + +def test_load_or_create_is_idempotent(tmp_path: Path) -> None: + store = McpTokenStore(tmp_path / "host_mcp_token.json") + first = store.load_or_create() + second = McpTokenStore(tmp_path / "host_mcp_token.json").load_or_create() + assert first.token == second.token + + +def test_load_or_create_writes_json_schema(tmp_path: Path) -> None: + path = tmp_path / "host_mcp_token.json" + McpTokenStore(path).load_or_create() + data = json.loads(path.read_text()) + assert set(data) == {"version", "token", "created_at"} + assert data["version"] == 1 + assert isinstance(data["token"], str) + # created_at is ISO 8601. + datetime.fromisoformat(data["created_at"]) + + +@pytest.mark.skipif(sys.platform == "win32", reason="POSIX perms only") +def test_load_or_create_sets_posix_permissions(tmp_path: Path) -> None: + path = tmp_path / "host_mcp_token.json" + McpTokenStore(path).load_or_create() + mode = stat.S_IMODE(os.fstat(os.open(path, os.O_RDONLY)).st_mode) + assert mode == 0o600 + + +def test_verify_accepts_correct_token(tmp_path: Path) -> None: + store = McpTokenStore(tmp_path / "host_mcp_token.json") + token = store.load_or_create() + assert store.verify(token.token) is True + + +def test_verify_rejects_wrong_token(tmp_path: Path) -> None: + store = McpTokenStore(tmp_path / "host_mcp_token.json") + store.load_or_create() + assert store.verify("wrong") is False + + +def test_load_or_create_raises_on_corrupt_json(tmp_path: Path) -> None: + path = tmp_path / "host_mcp_token.json" + path.write_text("{not valid json") + with pytest.raises(McpTokenStoreError): + McpTokenStore(path).load_or_create() + + +@pytest.mark.skipif(sys.platform == "win32", reason="POSIX chmod enforcement only") +def test_load_or_create_raises_on_unwritable_dir(tmp_path: Path) -> None: + unwritable = tmp_path / "ro" + unwritable.mkdir() + os.chmod(unwritable, 0o500) # r-x for owner + try: + with pytest.raises(McpTokenStoreError): + McpTokenStore(unwritable / "host_mcp_token.json").load_or_create() + finally: + os.chmod(unwritable, 0o700) # restore so cleanup works + + +@pytest.mark.skipif( + sys.platform == "win32", + reason="POSIX atomic-rename semantics only", +) +def test_load_or_create_concurrent_calls_do_not_corrupt( + tmp_path: Path, +) -> None: + """Two store instances racing to create: both end up reading the same token.""" + import threading + + path = tmp_path / "host_mcp_token.json" + results: list[McpToken] = [] + barrier = threading.Barrier(2) + + def worker() -> None: + barrier.wait() + store = McpTokenStore(path) + results.append(store.load_or_create()) + + threads = [threading.Thread(target=worker) for _ in range(2)] + for t in threads: + t.start() + for t in threads: + t.join() + assert len(results) == 2 + assert results[0].token == results[1].token