119 lines
3.4 KiB
Python
119 lines
3.4 KiB
Python
"""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
|