271 lines
9.4 KiB
Python
271 lines
9.4 KiB
Python
"""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 <token>``,
|
|
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)))
|