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()