mcp-server/build_retail_network.py

#!/usr/bin/env python3
"""
Parameterized retail network builder for GNS3.

Creates an HQ plus N store sites (1-10) with POS and WiFi endpoints wired into
an Internet/MPLS cloud. Designed to be driven by the MCP server tool
`build_network_from_description`.
"""

import json
import logging
import os
import time
from typing import Any, Dict, List, Optional, Tuple

import requests

logger = logging.getLogger(__name__)

GNS3_SERVER = os.getenv("GNS3_SERVER", "http://localhost:3080")
GNS3_API_BASE = f"{GNS3_SERVER.rstrip('/')}/v2"
PROJECT_NAME = os.getenv("GNS3_PROJECT_NAME", "overgrowth")
DEFAULT_SITE_COUNT = 3
IOSVL2_TEMPLATE_ID = "25dd7340-2e92-4e45-83de-f88077a24ceb"

STORE_SPACING = 220
STORE_CLOUD_Y = 200
STORE_SWITCH_Y = 100


def _clamp_site_count(site_count: int) -> int:
    try:
        value = int(site_count)
    except Exception:
        value = DEFAULT_SITE_COUNT
    return max(1, min(10, value))


def _ensure_project(session: requests.Session, api_base: str, name: str) -> Optional[Dict[str, Any]]:
    """Return an existing project by name or create one."""
    try:
        resp = session.get(f"{api_base}/projects")
        if resp.ok:
            project = next((p for p in resp.json() if p.get("name") == name), None)
            if project:
                return project
    except Exception:
        pass

    try:
        created = session.post(f"{api_base}/projects", json={"name": name})
        if created.status_code in (200, 201):
            return created.json()
    except Exception:
        pass
    return None


def _delete_all_nodes(session: requests.Session, api_base: str, project_id: str) -> int:
    """Remove all nodes from a project to ensure a clean rebuild."""
    deleted = 0
    resp = session.get(f"{api_base}/projects/{project_id}/nodes")
    if not resp.ok:
        return deleted
    for node in resp.json():
        try:
            session.delete(f"{api_base}/projects/{project_id}/nodes/{node['node_id']}")
            deleted += 1
        except Exception:
            continue
    time.sleep(1)
    return deleted


def _create_node(
    session: requests.Session, api_base: str, project_id: str, payload: Dict[str, Any], label: str
) -> Optional[Dict[str, Any]]:
    resp = session.post(f"{api_base}/projects/{project_id}/nodes", json=payload)
    if resp.status_code in (200, 201):
        node = resp.json()
        print(f"  ✓ {label} ({node.get('name')})")
        return node
    print(f"  ✗ Failed to create {label}: {resp.text}")
    return None


def create_switch(session: requests.Session, api_base: str, project_id: str, name: str, x: int, y: int):
    payload = {
        "name": name,
        "node_type": "qemu",
        "compute_id": "local",
        "template_id": IOSVL2_TEMPLATE_ID,
        "properties": {
            "qemu_path": "/usr/bin/qemu-system-x86_64",
            "platform": "x86_64",
            "adapters": 16,
            "ram": 1024,
            "hda_disk_image": "vios_l2-adventerprisek9-m.03.2017.qcow2",
        },
        "symbol": ":/symbols/multilayer_switch.svg",
        "x": x,
        "y": y,
    }
    return _create_node(session, api_base, project_id, payload, f"Switch {name}")


def create_cloud(session: requests.Session, api_base: str, project_id: str, name: str, x: int, y: int):
    payload = {
        "name": name,
        "node_type": "cloud",
        "compute_id": "local",
        "symbol": ":/symbols/cloud.svg",
        "x": x,
        "y": y,
    }
    return _create_node(session, api_base, project_id, payload, f"Cloud {name}")


def create_pc(
    session: requests.Session,
    api_base: str,
    project_id: str,
    name: str,
    x: int,
    y: int,
    symbol: str = ":/symbols/vpcs_guest.svg",
):
    payload = {
        "name": name,
        "node_type": "vpcs",
        "compute_id": "local",
        "symbol": symbol,
        "x": x,
        "y": y,
    }
    return _create_node(session, api_base, project_id, payload, f"VPCS {name}")


def create_link(
    session: requests.Session,
    api_base: str,
    project_id: str,
    node1_id: str,
    adapter1: int,
    port1: int,
    node2_id: str,
    adapter2: int,
    port2: int,
    desc: str,
) -> bool:
    payload = {
        "nodes": [
            {"node_id": node1_id, "adapter_number": adapter1, "port_number": port1},
            {"node_id": node2_id, "adapter_number": adapter2, "port_number": port2},
        ]
    }
    resp = session.post(f"{api_base}/projects/{project_id}/links", json=payload)
    success = resp.status_code in (200, 201)
    print(f"  {'✓' if success else '✗'} {desc}")
    return success


def _store_coordinates(index: int, total_sites: int) -> Tuple[int, int, int, int]:
    offset = (index - (total_sites + 1) / 2) * STORE_SPACING
    x_cloud = int(offset)
    x_switch = int(offset)
    return x_cloud, STORE_CLOUD_Y, x_switch, STORE_SWITCH_Y


def build_retail_network(
    gns3_api_base: str,
    project_name: str,
    site_count: int = DEFAULT_SITE_COUNT,
) -> Dict[str, Any]:
    """
    Build an HQ + N store retail topology in the given GNS3 project.
    """
    api_base = gns3_api_base.rstrip("/")
    session = requests.Session()
    site_count = _clamp_site_count(site_count)
    print("=" * 70)
    print(f"🌿 Building retail network for project '{project_name}' with {site_count} stores")

    project = _ensure_project(session, api_base, project_name)
    if not project:
        raise RuntimeError(f"Unable to create or locate project '{project_name}'")
    project_id = project.get("project_id") or project.get("id")

    _delete_all_nodes(session, api_base, project_id)

    nodes: Dict[str, Optional[Dict[str, Any]]] = {}
    errors: List[str] = []

    print("\n🏢 Building Headquarters...")
    nodes["hq_area"] = create_cloud(session, api_base, project_id, "🏢-HQ-Building", 0, -450)
    nodes["hq_core"] = create_switch(session, api_base, project_id, "SW-HQ-Core", 0, -300)
    nodes["hq_server"] = create_pc(
        session, api_base, project_id, "Server-DB", -150, -200, symbol=":/symbols/server.svg"
    )
    nodes["hq_pc"] = create_pc(
        session, api_base, project_id, "Office-Manager", 150, -200, symbol=":/symbols/computer.svg"
    )

    print("\n🌐 Building Internet/MPLS cloud...")
    nodes["internet"] = create_cloud(session, api_base, project_id, "☁️-Internet-MPLS", 0, -50)

    print("\n☕ Building store sites...")
    store_records: List[Dict[str, Any]] = []
    for idx in range(1, site_count + 1):
        x_cloud, y_cloud, x_switch, y_switch = _store_coordinates(idx, site_count)
        store_name = f"☕-Store-{idx}"
        switch_name = f"SW-Store-{idx}"
        pos_name = f"POS-Store-{idx}"
        wifi_name = f"WiFi-Store-{idx}"

        store_cloud = create_cloud(session, api_base, project_id, store_name, x_cloud, y_cloud)
        store_switch = create_switch(session, api_base, project_id, switch_name, x_switch, y_switch)
        store_pos = create_pc(session, api_base, project_id, pos_name, x_cloud - 80, y_cloud, ":/symbols/atm.svg")
        store_wifi = create_pc(
            session, api_base, project_id, wifi_name, x_cloud + 80, y_cloud, ":/symbols/wifi_antenna.svg"
        )
        store_records.append(
            {
                "index": idx,
                "cloud": store_cloud,
                "switch": store_switch,
                "pos": store_pos,
                "wifi": store_wifi,
            }
        )

    print("\n⏱️  Waiting for devices to initialize...")
    time.sleep(2)

    link_summaries: List[Dict[str, Any]] = []

    def add_link(src: Optional[Dict[str, Any]], adapter1: int, port1: int,
                 dst: Optional[Dict[str, Any]], adapter2: int, port2: int, desc: str):
        if not src or not dst:
            errors.append(f"Missing nodes for link: {desc}")
            link_summaries.append({"description": desc, "success": False})
            return
        success = create_link(
            session,
            api_base,
            project_id,
            src["node_id"],
            adapter1,
            port1,
            dst["node_id"],
            adapter2,
            port2,
            desc,
        )
        link_summaries.append(
            {
                "description": desc,
                "success": success,
                "from": src.get("name"),
                "to": dst.get("name"),
            }
        )
        if not success:
            errors.append(f"Link failed: {desc}")

    print("\n🔗 Connecting core and HQ devices...")
    add_link(nodes.get("hq_server"), 0, 0, nodes.get("hq_core"), 0, 0, "Server-DB → SW-HQ-Core")
    add_link(nodes.get("hq_pc"), 0, 0, nodes.get("hq_core"), 1, 0, "Office-Manager → SW-HQ-Core")
    add_link(nodes.get("hq_core"), 15, 0, nodes.get("internet"), 0, 0, "SW-HQ-Core → Internet/MPLS")

    print("\n🔗 Connecting stores to WAN and local devices...")
    for store in store_records:
        idx = store["index"]
        switch_name = store.get("switch", {}).get("name", f"SW-Store-{idx}")
        pos_name = store.get("pos", {}).get("name", f"POS-Store-{idx}")
        wifi_name = store.get("wifi", {}).get("name", f"WiFi-Store-{idx}")
        add_link(
            nodes.get("internet"),
            idx,
            0,
            store.get("switch"),
            15,
            0,
            f"Internet/MPLS → {switch_name}",
        )
        add_link(store.get("pos"), 0, 0, store.get("switch"), 0, 0, f"{pos_name} → {switch_name}")
        add_link(store.get("wifi"), 0, 0, store.get("switch"), 1, 0, f"{wifi_name} → {switch_name}")

    final_nodes: List[Dict[str, Any]] = []
    try:
        resp_nodes = session.get(f"{api_base}/projects/{project_id}/nodes")
        if resp_nodes.ok:
            final_nodes = resp_nodes.json()
    except Exception:
        pass

    store_clouds = [
        n.get("name", "") for n in final_nodes if n.get("node_type") == "cloud" and str(n.get("name", "")).startswith("☕-Store-")
    ]

    summary = {
        "success": len(errors) == 0,
        "project_name": project_name,
        "project_id": project_id,
        "site_count_requested": site_count,
        "site_count_built": len(store_clouds),
        "nodes": final_nodes,
        "links": link_summaries,
        "errors": errors,
    }

    print("\n✅ Retail network build complete")
    print(f"   Requested stores: {site_count} | Built stores: {len(store_clouds)}")
    logger.info("Retail builder created %d sites for project %s", len(store_clouds), project_name)
    return summary


def main():
    site_count_env = os.getenv("SITE_COUNT")
    try:
        site_count = int(site_count_env) if site_count_env else DEFAULT_SITE_COUNT
    except Exception:
        site_count = DEFAULT_SITE_COUNT

    try:
        result = build_retail_network(GNS3_API_BASE, PROJECT_NAME, site_count)
        print(json.dumps(result, indent=2))
    except Exception as e:
        print(json.dumps({"success": False, "error": str(e)}))


if __name__ == "__main__":
    main()