server.py•7.42 kB
"""Main MCP server for NetBox infrastructure access."""
import asyncio
import logging
import os
import sys
from typing import Any, Dict, Optional
from dotenv import load_dotenv
from mcp.server import Server
from mcp.server.stdio import stdio_server
from mcp.types import Tool, TextContent
from src.netbox_client import NetBoxClient
from src.state_confidence import StateConfidenceClient
from src.vault_client import VaultClient
from src.tools import hosts
from src.tools import ip_addresses
from src.tools import virtual_machines
from src.tools import vlans
# Load environment variables
load_dotenv()
# Configure logging
logging.basicConfig(
level=os.getenv("LOG_LEVEL", "INFO").upper(),
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger(__name__)
# Initialize global clients
vault_client: Optional[VaultClient] = None
netbox_client: Optional[NetBoxClient] = None
state_client: Optional[StateConfidenceClient] = None
def initialize_clients():
"""Initialize Vault, NetBox, and PostgreSQL clients."""
global vault_client, netbox_client, state_client
# Vault client
vault_addr = os.getenv("VAULT_ADDR", "http://localhost:8200")
vault_role_id = os.getenv("VAULT_ROLE_ID")
vault_secret_id = os.getenv("VAULT_SECRET_ID")
vault_mount_path = os.getenv("VAULT_MOUNT_PATH", "auth/approle")
if not vault_role_id or not vault_secret_id:
logger.warning(
"VAULT_ROLE_ID or VAULT_SECRET_ID not set. Vault authentication will fail."
)
vault_client = VaultClient(
vault_addr=vault_addr,
role_id=vault_role_id or "",
secret_id=vault_secret_id or "",
mount_path=vault_mount_path,
)
# Authenticate with Vault
if vault_client.authenticate():
logger.info("Vault client initialized successfully")
else:
logger.error("Failed to authenticate with Vault. NetBox access will fail.")
# NetBox client
netbox_url = os.getenv("NETBOX_URL", "http://localhost:8000")
vault_path = os.getenv("NETBOX_VAULT_PATH", "netbox/jit-tokens")
netbox_client = NetBoxClient(
netbox_url=netbox_url, vault_client=vault_client, vault_path=vault_path
)
logger.info(f"NetBox client initialized for {netbox_url}")
# State confidence client (optional)
postgres_host = os.getenv("POSTGRES_HOST")
postgres_port = os.getenv("POSTGRES_PORT", "5432")
postgres_db = os.getenv("POSTGRES_DB", "state")
postgres_user = os.getenv("POSTGRES_USER")
postgres_password = os.getenv("POSTGRES_PASSWORD", "")
postgres_sslmode = os.getenv("POSTGRES_SSLMODE", "prefer")
if postgres_host and postgres_user:
try:
state_client = StateConfidenceClient(
host=postgres_host,
port=int(postgres_port),
database=postgres_db,
user=postgres_user,
password=postgres_password,
sslmode=postgres_sslmode,
)
logger.info("State confidence client initialized successfully")
except Exception as e:
logger.warning(f"Failed to initialize state confidence client: {e}")
state_client = None
else:
logger.info(
"PostgreSQL configuration not provided. State confidence scores will not be available."
)
# Initialize clients at module load
initialize_clients()
# Create MCP server
server = Server(name="netbox-infrastructure")
@server.list_tools()
async def list_tools() -> list[Tool]:
"""List all available tools."""
if not netbox_client:
logger.error("NetBox client not initialized")
return []
tools = []
# Host tools
tools.extend(hosts.get_host_tools(netbox_client, state_client))
# VM tools
tools.extend(virtual_machines.get_vm_tools(netbox_client, state_client))
# IP tools
tools.extend(ip_addresses.get_ip_tools(netbox_client, state_client))
# VLAN tools
tools.extend(vlans.get_vlan_tools(netbox_client, state_client))
return tools
@server.call_tool()
async def call_tool(name: str, arguments: Dict[str, Any]) -> list[TextContent]:
"""Handle tool calls."""
if not netbox_client:
return [
TextContent(
type="text",
text="Error: NetBox client not initialized. Check configuration.",
)
]
try:
# Route to appropriate handler
if name == "list_hosts":
return await hosts.handle_list_hosts(arguments, netbox_client, state_client)
elif name == "get_host":
return await hosts.handle_get_host(arguments, netbox_client, state_client)
elif name == "search_hosts":
return await hosts.handle_search_hosts(arguments, netbox_client, state_client)
elif name == "list_vms":
return await virtual_machines.handle_list_vms(
arguments, netbox_client, state_client
)
elif name == "get_vm":
return await virtual_machines.handle_get_vm(
arguments, netbox_client, state_client
)
elif name == "list_vm_interfaces":
return await virtual_machines.handle_list_vm_interfaces(
arguments, netbox_client, state_client
)
elif name == "list_ips":
return await ip_addresses.handle_list_ips(
arguments, netbox_client, state_client
)
elif name == "get_ip":
return await ip_addresses.handle_get_ip(
arguments, netbox_client, state_client
)
elif name == "search_ips":
return await ip_addresses.handle_search_ips(
arguments, netbox_client, state_client
)
elif name == "list_vlans":
return await vlans.handle_list_vlans(arguments, netbox_client, state_client)
elif name == "get_vlan":
return await vlans.handle_get_vlan(arguments, netbox_client, state_client)
elif name == "list_vlan_ips":
return await vlans.handle_list_vlan_ips(
arguments, netbox_client, state_client
)
else:
return [
TextContent(
type="text",
text=f"Error: Unknown tool '{name}'",
)
]
except Exception as e:
logger.error(f"Error handling tool call '{name}': {e}", exc_info=True)
return [
TextContent(
type="text",
text=f"Error executing tool '{name}': {str(e)}",
)
]
async def main():
"""Main entry point for MCP server."""
logger.info("Starting NetBox Infrastructure MCP Server...")
logger.info(f"NetBox URL: {os.getenv('NETBOX_URL', 'http://localhost:8000')}")
logger.info(f"Vault Address: {os.getenv('VAULT_ADDR', 'http://localhost:8200')}")
async with stdio_server() as (read_stream, write_stream):
await server.run(
read_stream,
write_stream,
server.create_initialization_options(),
)
if __name__ == "__main__":
try:
asyncio.run(main())
except KeyboardInterrupt:
logger.info("Server stopped by user")
sys.exit(0)
except Exception as e:
logger.error(f"Fatal error: {e}", exc_info=True)
sys.exit(1)