ctfd-mcp/server.py

from __future__ import annotations

from contextlib import asynccontextmanager
from typing import Any

import anyio
from mcp.server.fastmcp import FastMCP

from config import ConfigError, load_config
from ctfd_client import (
    AuthError,
    CTFdClient,
    CTFdClientError,
    NotFoundError,
    RateLimitError,
)

_client: CTFdClient | None = None


def _format_error(exc: Exception) -> RuntimeError:
    if isinstance(exc, ConfigError):
        message = f"Configuration error: {exc}"
    elif isinstance(exc, AuthError):
        message = f"Auth failed: {exc}"
    elif isinstance(exc, NotFoundError):
        message = f"Not found: {exc}"
    elif isinstance(exc, RateLimitError):
        retry = (
            f" Retry-After={exc.retry_after}."
            if getattr(exc, "retry_after", None)
            else ""
        )
        message = f"Rate limited.{retry}".strip()
    elif isinstance(exc, CTFdClientError):
        message = f"CTFd API error: {exc}"
    else:
        message = f"Unexpected error: {exc}"
    return RuntimeError(message)


async def _close_client() -> None:
    global _client
    if _client is None:
        return
    try:
        await _client.aclose()
    finally:
        _client = None


@asynccontextmanager
async def _lifespan(_: FastMCP):
    try:
        yield
    finally:
        await _close_client()


mcp = FastMCP("ctfd-mcp", lifespan=_lifespan)


async def _get_client() -> CTFdClient:
    global _client
    if _client is None:
        config = load_config()
        _client = CTFdClient(config)
    return _client


@mcp.tool(
    description="List visible challenges. Optional filter by category and unsolved only."
)
async def list_challenges(category: str | None = None, only_unsolved: bool = False):
    client = await _get_client()
    try:
        return await client.list_challenges(
            category=category, only_unsolved=only_unsolved
        )
    except Exception as exc:  # noqa: BLE001 - map to user-friendly MCP error
        raise _format_error(exc)


@mcp.tool(description="Get challenge details (description, files, meta) by ID.")
async def challenge_details(challenge_id: int):
    client = await _get_client()
    try:
        return await client.get_challenge(challenge_id)
    except Exception as exc:  # noqa: BLE001 - map to user-friendly MCP error
        raise _format_error(exc)


@mcp.tool(description="Submit a flag for a challenge ID.")
async def submit_flag(challenge_id: int, flag: str):
    client = await _get_client()
    try:
        return await client.submit_flag(challenge_id, flag)
    except Exception as exc:  # noqa: BLE001 - map to user-friendly MCP error
        raise _format_error(exc)


@mcp.tool(
    description="Unified start: detects plugin (whale/ctfd-owl/k8s) and starts container."
)
async def start_container(challenge_id: int):
    client = await _get_client()
    try:
        return await client.start_container(challenge_id)
    except Exception as exc:  # noqa: BLE001
        raise _format_error(exc)


@mcp.tool(
    description="Unified stop: whale requires container_id; ctfd-owl/k8s require challenge_id."
)
async def stop_container(
    container_id: int | None = None, challenge_id: int | None = None
):
    client = await _get_client()
    try:
        return await client.stop_container(
            container_id=container_id, challenge_id=challenge_id
        )
    except Exception as exc:  # noqa: BLE001
        raise _format_error(exc)


@mcp.resource(
    "resource://ctfd/challenges/{challenge_id}",
    name="ctfd-challenge",
    title="CTFd challenge details",
    description="Challenge description plus attachment URLs.",
    mime_type="text/markdown",
)
async def challenge_resource(challenge_id: int):
    client = await _get_client()
    try:
        details = await client.get_challenge(challenge_id)
    except Exception as exc:  # noqa: BLE001
        raise _format_error(exc)
    return _challenge_markdown(details)


def _challenge_markdown(details: dict[str, Any]) -> str:
    """Render a chat-friendly markdown snapshot of a challenge."""
    lines: list[str] = []
    title = details.get("name") or f"Challenge {details.get('id')}"
    lines.append(f"# {title}")
    subtitle: list[str] = []
    if details.get("id") is not None:
        subtitle.append(f"ID: {details['id']}")
    if details.get("category"):
        subtitle.append(f"Category: {details['category']}")
    if details.get("value") is not None:
        subtitle.append(f"Points: {details['value']}")
    if details.get("solved") is not None:
        subtitle.append("Solved" if details["solved"] else "Unsolved")
    if subtitle:
        lines.append(" / ".join(subtitle))

    desc = (details.get("description_text") or details.get("description") or "").strip()
    if desc:
        lines.append("")
        lines.append("## Description")
        lines.append(desc)

    conn = details.get("connection_info")
    if conn:
        lines.append("")
        lines.append("## Connection")
        lines.append(str(conn))

    files = details.get("files") or []
    if files:
        lines.append("")
        lines.append("## Files")
        lines.extend(f"- {url}" for url in files)

    return "\n".join(lines).strip() or "No challenge details available."


def run(transport: str = "stdio") -> None:
    """Start the MCP server."""
    if transport == "stdio":
        anyio.run(_run_stdio_with_lifecycle)
        return
    mcp.run(transport=transport)


async def _run_stdio_with_lifecycle() -> None:
    async with _lifespan(mcp):
        await mcp.run_stdio_async()