feat(host-agent): add McpTokenStore for MCP bearer token

2026-07-21 13:55:24 +08:00
parent 29b9a8c39a
commit c7faee8da3
2 changed files with 225 additions and 0 deletions
@@ -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