MCP enhancements #1343

This commit is contained in:
Jokob @NetAlertX
2025-12-12 05:21:23 +00:00
parent a627cc6abe
commit aed7a91bf0
8 changed files with 1110 additions and 826 deletions

View File

@@ -1,4 +1,15 @@
#!/usr/bin/env python
"""
NetAlertX MCP (Model Context Protocol) Server Endpoint.
This module implements an MCP server that exposes NetAlertX API endpoints as tools
for AI assistants. It provides JSON-RPC over HTTP and Server-Sent Events (SSE)
for tool discovery and execution.
The server maps OpenAPI specifications to MCP tools, allowing AIs to list available
tools and call them with appropriate parameters. Tools include device management,
network scanning, event querying, and more.
"""
import threading
from flask import Blueprint, request, jsonify, Response, stream_with_context
@@ -14,11 +25,18 @@ import queue
mcp_bp = Blueprint('mcp', __name__)
tools_bp = Blueprint('tools', __name__)
# Global session management for MCP SSE connections
mcp_sessions = {}
mcp_sessions_lock = threading.Lock()
def check_auth():
"""
Check if the request has valid authorization.
Returns:
bool: True if the Authorization header matches the expected API token, False otherwise.
"""
token = request.headers.get("Authorization")
expected_token = f"Bearer {get_setting_value('API_TOKEN')}"
return token == expected_token
@@ -28,6 +46,15 @@ def check_auth():
# Specs
# --------------------------
def openapi_spec():
"""
Generate the OpenAPI specification for NetAlertX tools.
This function returns a JSON representation of the available API endpoints
that are exposed as MCP tools, including paths, methods, and operation IDs.
Returns:
flask.Response: A JSON response containing the OpenAPI spec.
"""
# Spec matching actual available routes for MCP tools
mylog("verbose", ["[MCP] OpenAPI spec requested"])
spec = {
@@ -35,17 +62,112 @@ def openapi_spec():
"info": {"title": "NetAlertX Tools", "version": "1.1.0"},
"servers": [{"url": "/"}],
"paths": {
"/devices/by-status": {"post": {"operationId": "list_devices"}},
"/device/{mac}": {"post": {"operationId": "get_device_info"}},
"/devices/search": {"post": {"operationId": "search_devices"}},
"/devices/latest": {"get": {"operationId": "get_latest_device"}},
"/nettools/trigger-scan": {"post": {"operationId": "trigger_scan"}},
"/device/open_ports": {"post": {"operationId": "get_open_ports"}},
"/devices/network/topology": {"get": {"operationId": "get_network_topology"}},
"/events/recent": {"get": {"operationId": "get_recent_alerts"}, "post": {"operationId": "get_recent_alerts"}},
"/events/last": {"get": {"operationId": "get_last_events"}, "post": {"operationId": "get_last_events"}},
"/device/{mac}/set-alias": {"post": {"operationId": "set_device_alias"}},
"/nettools/wakeonlan": {"post": {"operationId": "wol_wake_device"}}
"/devices/by-status": {
"post": {
"operationId": "list_devices",
"description": "List devices filtered by their online/offline status. "
"Accepts optional 'status' query parameter (online/offline)."
}
},
"/device/{mac}": {
"post": {
"operationId": "get_device_info",
"description": "Retrieve detailed information about a specific device by MAC address."
}
},
"/devices/search": {
"post": {
"operationId": "search_devices",
"description": "Search for devices based on various criteria like name, IP, etc. "
"Accepts JSON with 'query' field."
}
},
"/devices/latest": {
"get": {
"operationId": "get_latest_device",
"description": "Get information about the most recently seen device."
}
},
"/devices/favorite": {
"get": {
"operationId": "get_favorite_devices",
"description": "Get favorite devices."
}
},
"/nettools/trigger-scan": {
"post": {
"operationId": "trigger_scan",
"description": "Trigger a network scan to discover new devices. "
"Accepts optional 'type' parameter for scan type - needs to match an enabled plugin name (e.g., ARPSCAN, NMAPDEV, NMAP)."
}
},
"/device/open_ports": {
"post": {
"operationId": "get_open_ports",
"description": "Get a list of open ports for a specific device. "
"Accepts JSON with 'target' (IP or MAC address). Trigger NMAP scan if no previous ports found with the /nettools/trigger-scan endpoint."
}
},
"/devices/network/topology": {
"get": {
"operationId": "get_network_topology",
"description": "Retrieve the network topology information."
}
},
"/events/recent": {
"get": {
"operationId": "get_recent_alerts",
"description": "Get recent events/alerts from the system. Defaults to last 24 hours."
},
"post": {"operationId": "get_recent_alerts"}
},
"/events/last": {
"get": {
"operationId": "get_last_events",
"description": "Get the last 10 events logged in the system."
},
"post": {"operationId": "get_last_events"}
},
"/device/{mac}/set-alias": {
"post": {
"operationId": "set_device_alias",
"description": "Set or update the alias/name for a device. Accepts JSON with 'alias' field."
}
},
"/nettools/wakeonlan": {
"post": {
"operationId": "wol_wake_device",
"description": "Send a Wake-on-LAN packet to wake up a device. "
"Accepts JSON with 'devMac' or 'devLastIP'."
}
},
"/devices/export": {
"get": {
"operationId": "export_devices",
"description": "Export devices in CSV or JSON format. "
"Accepts optional 'format' query parameter (csv/json, defaults to csv)."
}
},
"/devices/import": {
"post": {
"operationId": "import_devices",
"description": "Import devices from CSV or JSON content. "
"Accepts JSON with 'content' field containing base64-encoded data, or multipart file upload."
}
},
"/devices/totals": {
"get": {
"operationId": "get_device_totals",
"description": "Get device statistics and counts."
}
},
"/nettools/traceroute": {
"post": {
"operationId": "traceroute",
"description": "Perform a traceroute to a target IP address. "
"Accepts JSON with 'devLastIP' field."
}
}
}
}
return jsonify(spec)
@@ -57,11 +179,20 @@ def openapi_spec():
# Sessions for SSE
_openapi_spec_cache = None
API_BASE_URL = f"http://localhost:{get_setting_value('GRAPHQL_PORT')}"
_openapi_spec_cache = None # Cached OpenAPI spec to avoid repeated generation
API_BASE_URL = f"http://localhost:{get_setting_value('GRAPHQL_PORT')}" # Base URL for internal API calls
def get_openapi_spec():
"""
Retrieve the cached OpenAPI specification for MCP tools.
This function caches the OpenAPI spec to avoid repeated generation.
If the cache is empty, it calls openapi_spec() to generate it.
Returns:
dict or None: The OpenAPI spec as a dictionary, or None if generation fails.
"""
global _openapi_spec_cache
if _openapi_spec_cache:
@@ -78,6 +209,15 @@ def get_openapi_spec():
def map_openapi_to_mcp_tools(spec):
"""
Convert an OpenAPI specification into MCP tool definitions.
Args:
spec (dict): The OpenAPI spec dictionary.
Returns:
list: A list of MCP tool dictionaries, each containing name, description, and inputSchema.
"""
tools = []
if not spec or 'paths' not in spec:
return tools
@@ -101,6 +241,18 @@ def map_openapi_to_mcp_tools(spec):
def process_mcp_request(data):
"""
Process an incoming MCP JSON-RPC request.
Handles various MCP methods like initialize, tools/list, tools/call, etc.
For tools/call, it maps the tool name to an API endpoint and makes the call.
Args:
data (dict): The JSON-RPC request data containing method, id, params, etc.
Returns:
dict or None: The JSON-RPC response, or None for notifications.
"""
method = data.get('method')
msg_id = data.get('id')
if method == 'initialize':
@@ -157,6 +309,15 @@ def process_mcp_request(data):
def mcp_messages():
"""
Handle MCP messages for a specific session via HTTP POST.
This endpoint processes JSON-RPC requests for an existing MCP session.
The session_id is passed as a query parameter.
Returns:
flask.Response: JSON response indicating acceptance or error.
"""
session_id = request.args.get('session_id')
if not session_id:
return jsonify({"error": "Missing session_id"}), 400
@@ -174,6 +335,16 @@ def mcp_messages():
def mcp_sse():
"""
Handle MCP Server-Sent Events (SSE) endpoint.
Supports both GET (for establishing SSE stream) and POST (for direct JSON-RPC).
For GET, creates a new session and streams responses.
For POST, processes the request directly and returns the response.
Returns:
flask.Response: SSE stream for GET, JSON response for POST.
"""
if request.method == 'POST':
try:
data = request.get_json(silent=True)