feat: basics
This commit was merged in pull request #1.
This commit is contained in:
@@ -0,0 +1,339 @@
|
||||
"""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 groups ─────────────────────────────────────────────────────────
|
||||
|
||||
TOOL_GROUPS: dict[str, set[str]] = {
|
||||
"read-only": {
|
||||
"get_emails",
|
||||
"get_email",
|
||||
"search_emails",
|
||||
"get_calendar_events",
|
||||
"get_event_links",
|
||||
"find_person",
|
||||
"get_folders",
|
||||
"check_session",
|
||||
"find_free_time",
|
||||
"find_meeting_time",
|
||||
"get_meeting_stats",
|
||||
"get_meeting_contacts",
|
||||
"download_event_attachments",
|
||||
"get_email_links",
|
||||
"download_attachments",
|
||||
"login",
|
||||
},
|
||||
"write-only": {
|
||||
"send_email",
|
||||
"reply_email",
|
||||
"forward_email",
|
||||
"delete_email",
|
||||
"move_email",
|
||||
"mark_email_read",
|
||||
"create_meeting",
|
||||
"update_meeting",
|
||||
"cancel_meeting",
|
||||
"respond_to_meeting",
|
||||
"create_folder",
|
||||
"rename_folder",
|
||||
"empty_folder",
|
||||
"delete_folder",
|
||||
"move_folder",
|
||||
},
|
||||
}
|
||||
|
||||
# ── Alias resolution ────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _resolve_aliases(names: list[str]) -> set[str]:
|
||||
resolved: set[str] = set()
|
||||
for name in names:
|
||||
if name in TOOL_GROUPS:
|
||||
resolved |= TOOL_GROUPS[name]
|
||||
else:
|
||||
resolved.add(name)
|
||||
return resolved
|
||||
|
||||
|
||||
# ── 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
|
||||
|
||||
# ── Group integrity check ───────────────────────────────────────────────
|
||||
all_grouped: set[str] = set().union(*TOOL_GROUPS.values())
|
||||
missing = canonical - all_grouped
|
||||
assert not missing, f"Tools not in any group: {sorted(missing)}"
|
||||
overlaps: set[str] = set()
|
||||
for i, (g1, s1) in enumerate(TOOL_GROUPS.items()):
|
||||
for g2, s2 in list(TOOL_GROUPS.items())[i + 1 :]:
|
||||
overlaps |= s1 & s2
|
||||
assert not overlaps, f"Tools in multiple groups: {sorted(overlaps)}"
|
||||
|
||||
allowed: set[str] = set(canonical)
|
||||
|
||||
# ── Resolve aliases ────────────────────────────────────────────────────
|
||||
if enabled:
|
||||
enabled = sorted(_resolve_aliases(enabled))
|
||||
if disabled:
|
||||
disabled = sorted(_resolve_aliases(disabled))
|
||||
|
||||
# ── 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)))
|
||||
Reference in New Issue
Block a user