"""Shared lifecycle utilities for the OWA MCP server. Centralises three concerns so that tool modules do not duplicate code: 1. ``create_owa_client`` — builds and returns an initialised :class:`OWAClient`. 2. ``require_client`` — fetches the client from the lifespan context; raises a clear :class:`RuntimeError` when it is ``None`` (e.g. OWA URL was not passed). 3. ``apply_tool_filters`` — reads CLI args and applies whitelist / blacklist filtering to the tool registry. 4. ``SessionKeepalive`` — background task that pings OWA every *n* seconds to extend the session lifetime and persist fresh cookies to disk. """ from __future__ import annotations import asyncio import logging from typing import TYPE_CHECKING from mcp.server.auth.provider import AccessToken, TokenVerifier from owa_mcp.owa_client import OWAClient, SessionExpiredError if TYPE_CHECKING: from mcp.server.fastmcp import FastMCP logger = logging.getLogger(__name__) # ── Client factory ─────────────────────────────────────────────────────── def create_owa_client(owa_url: str, cookie_file: str | None = None) -> OWAClient: """Create and return a fully configured :class:`OWAClient`. Args: owa_url: Base URL of the OWA instance (e.g. ``https://owa.example.com``). cookie_file: Path to the session-cookie file. ``None`` falls back to the default ``session-cookies.txt`` next to the package. Returns: A ready-to-use :class:`OWAClient`. Raises: ValueError: If *owa_url* is empty. """ owa_url = owa_url.strip().rstrip("/") if not owa_url: raise ValueError("OWA URL must not be empty.") return OWAClient(owa_url=owa_url, cookie_file=cookie_file) # ── Token verifier ──────────────────────────────────────────────────────── class HashTokenVerifier(TokenVerifier): """Bearer-token verifier backed by a SHA-256 hash. The server never stores the plaintext token — only its hex digest. Clients send the plaintext token as ``Authorization: Bearer ``, the server hashes it and compares with the stored hash. """ def __init__(self, token_hash: str) -> None: self._expected_hash = token_hash.strip().lower() async def verify_token(self, token: str) -> AccessToken | None: import hashlib if hashlib.sha256(token.encode()).hexdigest() == self._expected_hash: return AccessToken( token=token, client_id="mcp-client", scopes=[], ) return None # ── Lifespan helper ────────────────────────────────────────────────────── class _LazyClient: """Thin wrapper that resolves the OWAClient on first access. Stored as ``AppContext.client`` so that ``require_client`` can work with both a pre-built client and a ``None`` placeholder. """ def __init__(self, client: OWAClient | None) -> None: self._client = client @property def resolved(self) -> OWAClient: if self._client is None: raise RuntimeError( "OWA client is not initialised. " "Check --owa-url and --cookie-file arguments." ) return self._client def require_client(client: OWAClient | None) -> OWAClient: """Return *client* if it is not ``None``, otherwise raise ``RuntimeError``. Import this in tool modules and call it with ``ctx.client`` to get a consistent error message when OWA is not configured. """ if client is None: raise RuntimeError( "OWA client is not initialised. Pass --owa-url when starting the server." ) return client # ── Session keepalive ────────────────────────────────────────────────────── _DEFAULT_KEEPALIVE_INTERVAL = 300 # seconds (5 minutes) class SessionKeepalive: """Background task that periodically pings the OWA server to extend the session lifetime and persist fresh cookies to disk. On each successful ping the OWA server may return updated ``Set-Cookie`` headers (e.g. a new ``X-OWA-CANARY`` or a renewed ``multifactor`` JWT). The client saves those cookies to disk automatically. On failure the keepalive attempts to reload cookies from disk so that a fresh login performed by another process (or via the ``login`` tool) is picked up without restarting the server. """ def __init__( self, client: OWAClient, interval: int = _DEFAULT_KEEPALIVE_INTERVAL ) -> None: self._client = client self._interval = interval self._task: asyncio.Task | None = None async def start(self) -> None: """Start the background keepalive loop. Fires one ping immediately so the session is refreshed on server start-up, then waits ``interval`` seconds between subsequent pings. """ if self._task is not None: return self._task = asyncio.create_task(self._run(), name="session-keepalive") async def stop(self) -> None: """Cancel the background keepalive loop and wait for it to finish.""" if self._task: self._task.cancel() try: await self._task except asyncio.CancelledError: pass self._task = None async def _run(self) -> None: """Ping immediately, then every ``interval`` seconds.""" # First ping right away (on server start) await asyncio.to_thread(self._ping) while True: await asyncio.sleep(self._interval) await asyncio.to_thread(self._ping) def _ping(self) -> None: """Make a lightweight GetFolder request against the inbox. On success the OWAClient saves fresh cookies to disk automatically. On ``SessionExpiredError`` the on-disk cookies are re-loaded in case the user re-authenticated while the server was running. """ try: self._client.request( "GetFolder", { "__type": "GetFolderJsonRequest:#Exchange", "Header": { "__type": "JsonRequestHeaders:#Exchange", "RequestServerVersion": "Exchange2013", }, "Body": { "__type": "GetFolderRequest:#Exchange", "FolderShape": { "__type": "FolderResponseShape:#Exchange", "BaseShape": "IdOnly", }, "FolderIds": [ { "__type": "DistinguishedFolderId:#Exchange", "Id": "inbox", } ], }, }, ) logger.debug("Session keepalive ping succeeded") except SessionExpiredError: logger.warning( "Session expired during keepalive — reloading cookies from disk" ) try: self._client._ensure_loaded() except Exception as exc: logger.warning( "Failed to reload cookies after keepalive failure: %s", exc ) except Exception as exc: logger.debug("Keepalive ping failed (non-critical): %s", exc) # ── Tool filtering ─────────────────────────────────────────────────────── def apply_tool_filters( server: FastMCP, enabled: list[str] | None = None, disabled: list[str] | None = None, ) -> None: """Apply whitelist / blacklist filtering to *server*'s registered tools. Args: server: The :class:`FastMCP` instance whose tool registry is modified in-place. enabled: Optional list of tool names to **keep**. If ``None`` or empty all tools are eligible. Unknown names are ignored (warning logged). disabled: Optional list of tool names to **remove**. Applied *after* the whitelist. Unknown names are ignored (warning logged). """ canonical = {t.name for t in asyncio.run(server.list_tools())} if not canonical: return allowed: set[str] = set(canonical) # ── Whitelist ────────────────────────────────────────────────────────── if enabled: enabled_set = {n.strip() for n in enabled if n.strip()} allowed &= enabled_set unknown = enabled_set - canonical if unknown: logger.warning( "Unknown tool(s) in --enabled-tools, ignoring: %s", ", ".join(sorted(unknown)), ) # ── Blacklist ────────────────────────────────────────────────────────── if disabled: disabled_set = {n.strip() for n in disabled if n.strip()} allowed -= disabled_set unknown = disabled_set - canonical if unknown: logger.warning( "Unknown tool(s) in --disabled-tools, ignoring: %s", ", ".join(sorted(unknown)), ) # ── Remove ───────────────────────────────────────────────────────────── to_remove = sorted(canonical - allowed) if to_remove: logger.info("Removing %d tool(s): %s", len(to_remove), ", ".join(to_remove)) for name in to_remove: server.remove_tool(name) logger.info("Tools available: %s", ", ".join(sorted(allowed)))