vlans.py•7.53 kB
"""MCP tools for VLAN queries."""
import logging
from typing import Any, Dict, List, Optional
from mcp.types import Tool, TextContent
from src.netbox_client import NetBoxClient
from src.state_confidence import StateConfidenceClient
logger = logging.getLogger(__name__)
def get_vlan_tools(
netbox_client: NetBoxClient, state_client: Optional[StateConfidenceClient]
) -> List[Tool]:
"""Get all VLAN-related MCP tools.
Args:
netbox_client: NetBox API client
state_client: Optional state confidence client
Returns:
List of Tool definitions
"""
return [
Tool(
name="list_vlans",
description="List all VLANs from NetBox. Supports filtering by VLAN ID, name, or site.",
inputSchema={
"type": "object",
"properties": {
"vid": {
"type": "integer",
"description": "Filter by VLAN ID",
},
"name": {
"type": "string",
"description": "Filter by VLAN name (partial match supported)",
},
"site": {
"type": "string",
"description": "Filter by site slug",
},
"limit": {
"type": "integer",
"description": "Maximum number of results (default: 100)",
"default": 100,
},
"include_certainty": {
"type": "boolean",
"description": "Include state confidence scores in response",
"default": True,
},
},
},
),
Tool(
name="get_vlan",
description="Get detailed information about a specific VLAN by VLAN ID.",
inputSchema={
"type": "object",
"properties": {
"vlan_id": {
"type": "integer",
"description": "VLAN ID",
},
"include_certainty": {
"type": "boolean",
"description": "Include state confidence score in response",
"default": True,
},
},
"required": ["vlan_id"],
},
),
Tool(
name="list_vlan_ips",
description="List all IP addresses assigned to a specific VLAN.",
inputSchema={
"type": "object",
"properties": {
"vlan_id": {
"type": "integer",
"description": "VLAN ID",
},
"include_certainty": {
"type": "boolean",
"description": "Include state confidence scores in response",
"default": True,
},
},
"required": ["vlan_id"],
},
),
]
async def handle_list_vlans(
arguments: Dict[str, Any],
netbox_client: NetBoxClient,
state_client: Optional[StateConfidenceClient],
) -> List[TextContent]:
"""Handle list_vlans tool call.
Args:
arguments: Tool arguments
netbox_client: NetBox API client
state_client: Optional state confidence client
Returns:
List of TextContent with results
"""
vid = arguments.get("vid")
name = arguments.get("name")
site = arguments.get("site")
limit = arguments.get("limit", 100)
include_certainty = arguments.get("include_certainty", True)
vlans = netbox_client.list_vlans(vid=vid, name=name, site=site, limit=limit)
results = []
for vlan in vlans:
result = vlan.copy()
if include_certainty and state_client:
vlan_id = str(vlan.get("vid", ""))
certainty = state_client.get_certainty_score(
state_id=vlan_id, state_type="vlan"
)
if certainty:
result["certainty"] = certainty
results.append(result)
import json
return [
TextContent(
type="text",
text=f"Found {len(results)} VLAN(s):\n\n{json.dumps(results, indent=2)}",
)
]
async def handle_get_vlan(
arguments: Dict[str, Any],
netbox_client: NetBoxClient,
state_client: Optional[StateConfidenceClient],
) -> List[TextContent]:
"""Handle get_vlan tool call.
Args:
arguments: Tool arguments
netbox_client: NetBox API client
state_client: Optional state confidence client
Returns:
List of TextContent with result
"""
vlan_id = arguments.get("vlan_id")
if vlan_id is None:
return [
TextContent(
type="text",
text="Error: vlan_id parameter is required",
)
]
include_certainty = arguments.get("include_certainty", True)
vlan = netbox_client.get_vlan(vlan_id)
if not vlan:
return [
TextContent(
type="text",
text=f"VLAN {vlan_id} not found in NetBox",
)
]
result = vlan.copy()
if include_certainty and state_client:
certainty = state_client.get_certainty_score(
state_id=str(vlan_id), state_type="vlan"
)
if certainty:
result["certainty"] = certainty
import json
return [
TextContent(
type="text",
text=f"VLAN details:\n\n{json.dumps(result, indent=2)}",
)
]
async def handle_list_vlan_ips(
arguments: Dict[str, Any],
netbox_client: NetBoxClient,
state_client: Optional[StateConfidenceClient],
) -> List[TextContent]:
"""Handle list_vlan_ips tool call.
Args:
arguments: Tool arguments
netbox_client: NetBox API client
state_client: Optional state confidence client
Returns:
List of TextContent with results
"""
vlan_id = arguments.get("vlan_id")
if vlan_id is None:
return [
TextContent(
type="text",
text="Error: vlan_id parameter is required",
)
]
include_certainty = arguments.get("include_certainty", True)
ip_addresses = netbox_client.list_vlan_ip_addresses(vlan_id)
if ip_addresses is None:
return [
TextContent(
type="text",
text=f"VLAN {vlan_id} not found in NetBox",
)
]
results = []
for ip_addr in ip_addresses:
result = ip_addr.copy()
if include_certainty and state_client:
ip_addr_str = ip_addr.get("address", "")
# Remove CIDR notation for state_id lookup
state_id = ip_addr_str.split("/")[0] if "/" in ip_addr_str else ip_addr_str
certainty = state_client.get_certainty_score(
state_id=state_id, state_type="ip_address"
)
if certainty:
result["certainty"] = certainty
results.append(result)
import json
return [
TextContent(
type="text",
text=f"Found {len(results)} IP address(es) in VLAN {vlan_id}:\n\n{json.dumps(results, indent=2)}",
)
]