feat: MCP server over SSE with OIDC auth
All checks were successful
CI / lint-and-test (push) Successful in 36s
All checks were successful
CI / lint-and-test (push) Successful in 36s
- Extract shared MCP tool handlers to mcp_common.py - mcp_server.py now uses shared handlers (stdio transport for local dev) - New routes/mcp.py: SSE transport behind existing OIDC Bearer auth - Mount MCP ASGI app at /mcp in main.py when AI_FEATURES_ENABLED - /mcp/sse -> establishes SSE stream (requires valid token when auth enabled) - /mcp/messages/ -> receives MCP client messages - Update README with SSE MCP docs - Add tests for mount existence, auth, and message routing
This commit is contained in:
@@ -1,9 +1,9 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
AOC MCP Server
|
||||
AOC MCP Server — stdio transport
|
||||
|
||||
Standalone MCP server that exposes audit log search tools for Claude Desktop,
|
||||
Cursor, and other MCP clients.
|
||||
Standalone MCP server for local use (Claude Desktop, Cursor, etc.).
|
||||
For the HTTP/SSE version (production, behind auth), see routes/mcp.py.
|
||||
|
||||
Usage:
|
||||
python mcp_server.py
|
||||
@@ -21,65 +21,28 @@ Claude Desktop config (~/.config/claude/claude_desktop_config.json):
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from datetime import UTC, datetime, timedelta
|
||||
|
||||
# Ensure backend modules are importable
|
||||
# Ensure backend modules are importable when run standalone
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
|
||||
from database import events_collection
|
||||
from mcp.server import Server
|
||||
from mcp.server.stdio import stdio_server
|
||||
from mcp.types import TextContent, Tool
|
||||
from mcp_common import (
|
||||
ASK_SCHEMA,
|
||||
GET_EVENT_SCHEMA,
|
||||
GET_SUMMARY_SCHEMA,
|
||||
SEARCH_EVENTS_SCHEMA,
|
||||
handle_ask,
|
||||
handle_get_event,
|
||||
handle_get_summary,
|
||||
handle_search_events,
|
||||
)
|
||||
|
||||
app = Server("aoc")
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tool definitions
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_SEARCH_EVENTS_SCHEMA = {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"entity": {"type": "string", "description": "Device name, user UPN, or email to search for"},
|
||||
"services": {
|
||||
"type": "array",
|
||||
"items": {"type": "string"},
|
||||
"description": "Filter by service (e.g. Intune, Directory, Exchange)",
|
||||
},
|
||||
"operation": {"type": "string", "description": "Filter by operation name"},
|
||||
"result": {"type": "string", "description": "Filter by result (success, failure)"},
|
||||
"days": {"type": "integer", "description": "Number of days to look back (default 7)"},
|
||||
"limit": {"type": "integer", "description": "Max events to return (default 20)"},
|
||||
},
|
||||
}
|
||||
|
||||
_GET_EVENT_SCHEMA = {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"event_id": {"type": "string", "description": "The event ID to retrieve"},
|
||||
},
|
||||
"required": ["event_id"],
|
||||
}
|
||||
|
||||
_GET_SUMMARY_SCHEMA = {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"days": {"type": "integer", "description": "Number of days to summarise (default 7)"},
|
||||
},
|
||||
}
|
||||
|
||||
_ASK_SCHEMA = {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"question": {"type": "string", "description": "Natural language question about audit logs"},
|
||||
"days": {"type": "integer", "description": "Number of days to look back (default 7)"},
|
||||
},
|
||||
"required": ["question"],
|
||||
}
|
||||
|
||||
|
||||
@app.list_tools()
|
||||
async def list_tools() -> list[Tool]:
|
||||
@@ -87,186 +50,35 @@ async def list_tools() -> list[Tool]:
|
||||
Tool(
|
||||
name="search_events",
|
||||
description="Search audit events by entity, service, operation, or result.",
|
||||
inputSchema=_SEARCH_EVENTS_SCHEMA,
|
||||
inputSchema=SEARCH_EVENTS_SCHEMA,
|
||||
),
|
||||
Tool(name="get_event", description="Retrieve a single audit event by its ID.", inputSchema=_GET_EVENT_SCHEMA),
|
||||
Tool(name="get_event", description="Retrieve a single audit event by its ID.", inputSchema=GET_EVENT_SCHEMA),
|
||||
Tool(
|
||||
name="get_summary",
|
||||
description="Get an aggregated summary of audit activity for the last N days.",
|
||||
inputSchema=_GET_SUMMARY_SCHEMA,
|
||||
inputSchema=GET_SUMMARY_SCHEMA,
|
||||
),
|
||||
Tool(
|
||||
name="ask",
|
||||
description="Ask a natural language question about audit logs. Returns a narrative answer.",
|
||||
inputSchema=_ASK_SCHEMA,
|
||||
inputSchema=ASK_SCHEMA,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tool handlers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@app.call_tool()
|
||||
async def call_tool(name: str, arguments: dict) -> list[TextContent]:
|
||||
if name == "search_events":
|
||||
return await _handle_search_events(arguments)
|
||||
return await handle_search_events(arguments)
|
||||
if name == "get_event":
|
||||
return await _handle_get_event(arguments)
|
||||
return await handle_get_event(arguments)
|
||||
if name == "get_summary":
|
||||
return await _handle_get_summary(arguments)
|
||||
return await handle_get_summary(arguments)
|
||||
if name == "ask":
|
||||
return await _handle_ask(arguments)
|
||||
return await handle_ask(arguments)
|
||||
raise ValueError(f"Unknown tool: {name}")
|
||||
|
||||
|
||||
async def _handle_search_events(arguments: dict) -> list[TextContent]:
|
||||
days = arguments.get("days", 7)
|
||||
limit = min(arguments.get("limit", 20), 100)
|
||||
since = (datetime.now(UTC) - timedelta(days=days)).isoformat().replace("+00:00", "Z")
|
||||
|
||||
filters = [{"timestamp": {"$gte": since}}]
|
||||
|
||||
services = arguments.get("services")
|
||||
if services:
|
||||
filters.append({"service": {"$in": services}})
|
||||
|
||||
operation = arguments.get("operation")
|
||||
if operation:
|
||||
filters.append({"operation": {"$regex": operation, "$options": "i"}})
|
||||
|
||||
result = arguments.get("result")
|
||||
if result:
|
||||
filters.append({"result": {"$regex": result, "$options": "i"}})
|
||||
|
||||
entity = arguments.get("entity")
|
||||
if entity:
|
||||
entity_safe = entity.replace(".", "\\.").replace("(", "\\(").replace(")", "\\)")
|
||||
filters.append(
|
||||
{
|
||||
"$or": [
|
||||
{"target_displays": {"$elemMatch": {"$regex": entity_safe, "$options": "i"}}},
|
||||
{"actor_display": {"$regex": entity_safe, "$options": "i"}},
|
||||
{"actor_upn": {"$regex": entity_safe, "$options": "i"}},
|
||||
{"raw_text": {"$regex": entity_safe, "$options": "i"}},
|
||||
]
|
||||
}
|
||||
)
|
||||
|
||||
query = {"$and": filters}
|
||||
cursor = events_collection.find(query).sort("timestamp", -1).limit(limit)
|
||||
events = list(cursor)
|
||||
|
||||
if not events:
|
||||
return [TextContent(type="text", text="No matching events found.")]
|
||||
|
||||
lines = [f"Found {len(events)} event(s):\n"]
|
||||
for e in events:
|
||||
ts = e.get("timestamp", "?")[:16].replace("T", " ")
|
||||
svc = e.get("service", "?")
|
||||
op = e.get("operation", "?")
|
||||
actor = e.get("actor_display", "?")
|
||||
result_str = e.get("result", "?")
|
||||
lines.append(f"{ts} | {svc} | {op} | {actor} | {result_str}")
|
||||
|
||||
return [TextContent(type="text", text="\n".join(lines))]
|
||||
|
||||
|
||||
async def _handle_get_event(arguments: dict) -> list[TextContent]:
|
||||
event_id = arguments["event_id"]
|
||||
event = events_collection.find_one({"id": event_id})
|
||||
if not event:
|
||||
return [TextContent(type="text", text=f"Event {event_id} not found.")]
|
||||
event.pop("_id", None)
|
||||
return [TextContent(type="text", text=json.dumps(event, indent=2, default=str))]
|
||||
|
||||
|
||||
async def _handle_get_summary(arguments: dict) -> list[TextContent]:
|
||||
days = arguments.get("days", 7)
|
||||
since = (datetime.now(UTC) - timedelta(days=days)).isoformat().replace("+00:00", "Z")
|
||||
query = {"timestamp": {"$gte": since}}
|
||||
|
||||
total = events_collection.count_documents(query)
|
||||
if total == 0:
|
||||
return [TextContent(type="text", text="No events in the specified period.")]
|
||||
|
||||
# Aggregation pipelines
|
||||
svc_pipeline = [
|
||||
{"$match": query},
|
||||
{"$group": {"_id": "$service", "count": {"$sum": 1}}},
|
||||
{"$sort": {"count": -1}},
|
||||
{"$limit": 10},
|
||||
]
|
||||
op_pipeline = [
|
||||
{"$match": query},
|
||||
{"$group": {"_id": "$operation", "count": {"$sum": 1}}},
|
||||
{"$sort": {"count": -1}},
|
||||
{"$limit": 10},
|
||||
]
|
||||
result_pipeline = [
|
||||
{"$match": query},
|
||||
{"$group": {"_id": "$result", "count": {"$sum": 1}}},
|
||||
{"$sort": {"count": -1}},
|
||||
]
|
||||
actor_pipeline = [
|
||||
{"$match": query},
|
||||
{"$group": {"_id": "$actor_display", "count": {"$sum": 1}}},
|
||||
{"$sort": {"count": -1}},
|
||||
{"$limit": 10},
|
||||
]
|
||||
|
||||
svc_counts = list(events_collection.aggregate(svc_pipeline))
|
||||
op_counts = list(events_collection.aggregate(op_pipeline))
|
||||
result_counts = list(events_collection.aggregate(result_pipeline))
|
||||
actor_counts = list(events_collection.aggregate(actor_pipeline))
|
||||
|
||||
lines = [f"Summary for the last {days} days ({total} total events)\n"]
|
||||
|
||||
lines.append("By service:")
|
||||
for row in svc_counts:
|
||||
lines.append(f" {row['_id'] or 'Unknown'}: {row['count']}")
|
||||
|
||||
lines.append("\nBy action:")
|
||||
for row in op_counts:
|
||||
lines.append(f" {row['_id'] or 'Unknown'}: {row['count']}")
|
||||
|
||||
lines.append("\nBy result:")
|
||||
for row in result_counts:
|
||||
lines.append(f" {row['_id'] or 'Unknown'}: {row['count']}")
|
||||
|
||||
lines.append("\nTop actors:")
|
||||
for row in actor_counts:
|
||||
lines.append(f" {row['_id'] or 'Unknown'}: {row['count']}")
|
||||
|
||||
return [TextContent(type="text", text="\n".join(lines))]
|
||||
|
||||
|
||||
async def _handle_ask(arguments: dict) -> list[TextContent]:
|
||||
"""For now, the MCP 'ask' tool returns a helpful message directing the user to the web UI,
|
||||
since the full NLQ pipeline requires LLM configuration that may not be available in the MCP context."""
|
||||
question = arguments["question"]
|
||||
days = arguments.get("days", 7)
|
||||
|
||||
# Perform a search to give the user something useful immediately
|
||||
result = await _handle_search_events({"entity": "", "days": days, "limit": 50})
|
||||
base_text = result[0].text if result else ""
|
||||
|
||||
text = (
|
||||
f"You asked: '{question}'\n\n"
|
||||
f"Here are the most recent {min(50, base_text.count(chr(10)) - 1)} events from the last {days} days:\n\n"
|
||||
f"{base_text}\n\n"
|
||||
f"Tip: Use the 'search_events' tool with specific filters (services, operation, result) "
|
||||
f"to narrow down the dataset before asking follow-up questions."
|
||||
)
|
||||
return [TextContent(type="text", text=text)]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Entry point
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def main():
|
||||
async with stdio_server() as (read_stream, write_stream):
|
||||
await app.run(read_stream, write_stream, app.create_initialization_options())
|
||||
|
||||
Reference in New Issue
Block a user