owa-mcp/owa_mcp/server.py

"""OWA MCP Server.

Exposes OWA email, calendar, directory, and availability tools via MCP.
Uses FastMCP with a lifespan context manager to share a single OWAClient
instance across all tool invocations.

CLI arguments:
    --transport           Transport: stdio (default), sse, streamable-http
    --owa-url             Base URL of the OWA instance (required)
    --cookie-file         Path to session cookies file (default: session-cookies.txt)
    --token-hash          SHA-256 hash of the bearer token (required for HTTP)
    --port                HTTP port (default 8000)
    --host                HTTP host (default 127.0.0.1)
    --enabled-tools       Comma-separated whitelist of tool names
    --disabled-tools      Comma-separated blacklist of tool names
"""

from __future__ import annotations

import argparse
import logging
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from dataclasses import dataclass

from mcp.server.auth.provider import AccessToken, TokenVerifier
from mcp.server.auth.settings import AuthSettings
from mcp.server.fastmcp import FastMCP

from owa_mcp.owa_client import OWAClient
from owa_mcp.server_lifecycle import (
    HashTokenVerifier,
    SessionKeepalive,
    apply_tool_filters,
    create_owa_client,
)

logger = logging.getLogger(__name__)


@dataclass
class AppContext:
    """Shared application state available to all tools via lifespan context."""
    client: OWAClient | None = None


@asynccontextmanager
async def app_lifespan(server: FastMCP) -> AsyncIterator[AppContext]:
    """Initialise the OWA client on first connection, then keep the session alive.

    A :class:`~owa_mcp.server_lifecycle.SessionKeepalive` task starts
    immediately (pings inbox, extends the session, persists fresh cookies) and
    repeats every 5 minutes until the server shuts down.
    """
    client = create_owa_client(_lifespan_owa_url, _lifespan_cookie_file)
    try:
        client._ensure_loaded()
    except Exception:
        pass  # no cookies yet — login tool will handle this

    keepalive = SessionKeepalive(client)
    await keepalive.start()

    try:
        yield AppContext(client=client)
    finally:
        await keepalive.stop()


# ── Temporary storage for lifespan args (set by main() before mcp.run()) ──
_lifespan_owa_url: str = ""
_lifespan_cookie_file: str | None = None


# ── Module-level MCP server (imported by all tool modules) ────────────────
mcp: FastMCP = FastMCP("owa", lifespan=app_lifespan)


def _parse_args(argv=None) -> argparse.Namespace:
    """Parse command-line arguments."""
    parser = argparse.ArgumentParser(
        description="OWA MCP Server — OWA/Exchange integration via MCP",
    )
    parser.add_argument(
        "--transport",
        choices=["stdio", "sse", "streamable-http"],
        default="stdio",
        help="Transport protocol (default: stdio)",
    )
    parser.add_argument(
        "--owa-url",
        required=True,
        help="Base URL of the OWA instance (e.g. https://owa.example.com)",
    )
    parser.add_argument(
        "--cookie-file",
        default=None,
        help="Path to session cookies file (default: session-cookies.txt)",
    )
    parser.add_argument(
        "--token-hash",
        default=None,
        help=(
            "SHA-256 hash of the bearer token for HTTP transport auth. "
            "Compute with: echo -n 'my-token' | sha256sum"
        ),
    )
    parser.add_argument(
        "--port",
        type=int,
        default=8000,
        help="Port for HTTP/SSE transport (default: 8000)",
    )
    parser.add_argument(
        "--host",
        default="127.0.0.1",
        help="Host for HTTP/SSE transport (default: 127.0.0.1)",
    )
    parser.add_argument(
        "--enabled-tools",
        default=None,
        help="Comma-separated whitelist of tool names to expose",
    )
    parser.add_argument(
        "--disabled-tools",
        default=None,
        help="Comma-separated blacklist of tool names to hide",
    )
    return parser.parse_args(argv)


def _configure_auth(
    mcp_server: FastMCP,
    token_hash: str | None,
) -> None:
    """Apply token-hash auth to *mcp_server* if *token_hash* is provided."""
    if not token_hash:
        return

    token_verifier: TokenVerifier = HashTokenVerifier(token_hash)
    mcp_server.settings.auth = AuthSettings(  # type: ignore[misc]
        issuer_url="http://localhost",  # type: ignore[arg-type]
        resource_server_url="http://localhost",  # type: ignore[arg-type]
    )
    mcp_server._token_verifier = token_verifier  # noqa: SLF001


def main(argv=None) -> None:
    """Entry point: parse CLI args, configure server, run transport."""
    global _lifespan_owa_url, _lifespan_cookie_file  # noqa: PLW0603

    args = _parse_args(argv)

    # ── Validate --owa-url ─────────────────────────────────────────────────
    owa_url = args.owa_url.strip().rstrip("/")
    if not owa_url:
        print(
            "ERROR: --owa-url must not be empty.",
            file=__import__("sys").stderr,
        )
        raise SystemExit(1)

    # ── Stash lifespan args (consumed by app_lifespan) ─────────────────────
    _lifespan_owa_url = owa_url
    _lifespan_cookie_file = args.cookie_file

    # ── Validate --token-hash for non-stdio transports ─────────────────────
    if args.transport != "stdio" and not args.token_hash:
        print(
            "ERROR: --token-hash is required when using HTTP transport.",
            file=__import__("sys").stderr,
        )
        print(
            "  Compute hash: echo -n 'my-token' | sha256sum",
            file=__import__("sys").stderr,
        )
        raise SystemExit(1)

    # ── Register tools (imports populate module-level `mcp`) ───────────────
    import owa_mcp.tools.email          # noqa: E402, F401
    import owa_mcp.tools.calendar       # noqa: E402, F401
    import owa_mcp.tools.people         # noqa: E402, F401
    import owa_mcp.tools.folders        # noqa: E402, F401
    import owa_mcp.tools.availability   # noqa: E402, F401
    import owa_mcp.tools.analytics      # noqa: E402, F401
    import owa_mcp.tools.auth           # noqa: E402, F401

    # ── Auth for HTTP transports ───────────────────────────────────────────
    _configure_auth(mcp, args.token_hash)

    # ── Tool filtering ─────────────────────────────────────────────────────
    enabled: list[str] | None = None
    if args.enabled_tools is not None:
        enabled = [t.strip() for t in args.enabled_tools.split(",") if t.strip()]

    disabled: list[str] | None = None
    if args.disabled_tools is not None:
        disabled = [t.strip() for t in args.disabled_tools.split(",") if t.strip()]

    if enabled or disabled:
        apply_tool_filters(mcp, enabled=enabled, disabled=disabled)

    # ── Host / port (already passed to FastMCP at module level; update if
    #   different from defaults — useful when server is re-invoked with
    #   different args in tests) ─────────────────────────────────────────────
    mcp.settings.host = args.host
    mcp.settings.port = args.port

    # ── Run transport ─────────────────────────────────────────────────────
    match args.transport:
        case "stdio":
            mcp.run()
        case "sse":
            mcp.run(transport="sse")
        case "streamable-http":
            mcp.run(transport="streamable-http")


if __name__ == "__main__":
    main()