virtual_machines.py•7.99 kB
"""MCP tools for virtual machine/container 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_vm_tools(
netbox_client: NetBoxClient, state_client: Optional[StateConfidenceClient]
) -> List[Tool]:
"""Get all VM-related MCP tools.
Args:
netbox_client: NetBox API client
state_client: Optional state confidence client
Returns:
List of Tool definitions
"""
return [
Tool(
name="list_vms",
description="List all virtual machines and containers from NetBox. Supports filtering by name, role, or API IP address.",
inputSchema={
"type": "object",
"properties": {
"name": {
"type": "string",
"description": "Filter by VM name (partial match supported)",
},
"role": {
"type": "string",
"description": "Filter by VM role slug",
},
"primary_ip": {
"type": "string",
"description": "Filter by primary IP address (API IP)",
},
"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_vm",
description="Get detailed information about a specific virtual machine or container by name.",
inputSchema={
"type": "object",
"properties": {
"hostname": {
"type": "string",
"description": "VM/container hostname",
},
"include_certainty": {
"type": "boolean",
"description": "Include state confidence score in response",
"default": True,
},
},
"required": ["hostname"],
},
),
Tool(
name="list_vm_interfaces",
description="List network interfaces for a virtual machine or container, including assigned IP addresses.",
inputSchema={
"type": "object",
"properties": {
"vm_name": {
"type": "string",
"description": "VM/container hostname",
},
"include_certainty": {
"type": "boolean",
"description": "Include state confidence scores in response",
"default": True,
},
},
"required": ["vm_name"],
},
),
]
async def handle_list_vms(
arguments: Dict[str, Any],
netbox_client: NetBoxClient,
state_client: Optional[StateConfidenceClient],
) -> List[TextContent]:
"""Handle list_vms tool call.
Args:
arguments: Tool arguments
netbox_client: NetBox API client
state_client: Optional state confidence client
Returns:
List of TextContent with results
"""
name = arguments.get("name")
role = arguments.get("role")
primary_ip = arguments.get("primary_ip")
limit = arguments.get("limit", 100)
include_certainty = arguments.get("include_certainty", True)
vms = netbox_client.list_virtual_machines(
name=name, role=role, primary_ip=primary_ip, limit=limit
)
results = []
for vm in vms:
result = vm.copy()
if include_certainty and state_client:
certainty = state_client.get_certainty_score(
state_id=vm.get("name", ""), state_type="vm"
)
if certainty:
result["certainty"] = certainty
results.append(result)
import json
return [
TextContent(
type="text",
text=f"Found {len(results)} VM/container(s):\n\n{json.dumps(results, indent=2)}",
)
]
async def handle_get_vm(
arguments: Dict[str, Any],
netbox_client: NetBoxClient,
state_client: Optional[StateConfidenceClient],
) -> List[TextContent]:
"""Handle get_vm tool call.
Args:
arguments: Tool arguments
netbox_client: NetBox API client
state_client: Optional state confidence client
Returns:
List of TextContent with result
"""
hostname = arguments.get("hostname")
if not hostname:
return [
TextContent(
type="text",
text="Error: hostname parameter is required",
)
]
include_certainty = arguments.get("include_certainty", True)
vm = netbox_client.get_virtual_machine(hostname)
if not vm:
return [
TextContent(
type="text",
text=f"VM/container '{hostname}' not found in NetBox",
)
]
result = vm.copy()
if include_certainty and state_client:
certainty = state_client.get_certainty_score(
state_id=hostname, state_type="vm"
)
if certainty:
result["certainty"] = certainty
import json
return [
TextContent(
type="text",
text=f"VM/container details:\n\n{json.dumps(result, indent=2)}",
)
]
async def handle_list_vm_interfaces(
arguments: Dict[str, Any],
netbox_client: NetBoxClient,
state_client: Optional[StateConfidenceClient],
) -> List[TextContent]:
"""Handle list_vm_interfaces tool call.
Args:
arguments: Tool arguments
netbox_client: NetBox API client
state_client: Optional state confidence client
Returns:
List of TextContent with results
"""
vm_name = arguments.get("vm_name")
if not vm_name:
return [
TextContent(
type="text",
text="Error: vm_name parameter is required",
)
]
include_certainty = arguments.get("include_certainty", True)
interfaces = netbox_client.list_vm_interfaces(vm_name)
if not interfaces:
return [
TextContent(
type="text",
text=f"VM/container '{vm_name}' not found or has no interfaces",
)
]
results = []
for interface in interfaces:
result = interface.copy()
if include_certainty and state_client:
# Get certainty for the VM itself
vm_certainty = state_client.get_certainty_score(
state_id=vm_name, state_type="vm"
)
if vm_certainty:
result["vm_certainty"] = vm_certainty
# Get certainty for IP addresses
ip_addresses = result.get("ip_addresses", [])
for ip_addr in ip_addresses:
ip_certainty = state_client.get_certainty_score(
state_id=ip_addr.get("address", ""), state_type="ip_address"
)
if ip_certainty:
ip_addr["certainty"] = ip_certainty
results.append(result)
import json
return [
TextContent(
type="text",
text=f"Found {len(results)} interface(s) for '{vm_name}':\n\n{json.dumps(results, indent=2)}",
)
]