gns3-server/gns3server/agent/gns3_copilot/gns3_client/gns3_topology_reader.py

# SPDX-License-Identifier: GPL-3.0-or-later
#
# GNS3-Copilot - AI-powered Network Lab Assistant for GNS3
#
# This file is part of GNS3-Copilot project.
#
# GNS3-Copilot is free software: you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the
# Free Software Foundation, either version 3 of the License, or (at your
# option) any later version.
#
# GNS3-Copilot is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
# or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License
# for more details.
#
# You should have received a copy of the GNU General Public License
# along with GNS3-Copilot. If not, see <https://www.gnu.org/licenses/>.
#
# Copyright (C) 2025 Yue Guobin (岳国宾)
# Author: Yue Guobin (岳国宾)
#
# Project Home: https://github.com/yueguobin/gns3-copilot
#

"""
GNS3 Topology Reader Tool

This module provides a LangChain BaseTool to retrieve the topology of a
specific GNS3 project by project ID. Returns nodes, links, and project
metadata.

⚠️ WARNING: This module is shared with the MCP (Model Context Protocol) service.
GNS3TopologyTool._run() is called by MCP device config handlers.
The jwt_token/url parameters were added for MCP compatibility.
Modifications must be tested with BOTH gns3-copilot AND MCP.
"""

import copy
import logging
from pprint import pprint
from typing import Any

from langchain.tools import BaseTool

from gns3server.agent.gns3_copilot.gns3_client import Project
from gns3server.agent.gns3_copilot.gns3_client import get_gns3_connector

# Configure logging
logger = logging.getLogger(__name__)


# Define LangChain tool class
class GNS3TopologyTool(BaseTool):
    """LangChain tool for retrieving GNS3 project topology information."""

    name: str = "gns3_topology_reader"
    description: str = """
    Retrieves the topology of a GNS3 project including nodes and links.

    Input: `project_id` (str, required): UUID of the GNS3 project.

    Output: Dictionary with:
    - `project_id`, `name`, `status`: Project metadata
    - `nodes`: Dict of node details (node_id, name, ports, console_port,
               type, etc.)
    - `links`: List of link connections

    Use this to understand network structure before making changes.
    """

    def _run(
        self,
        tool_input: Any = None,
        run_manager: Any = None,
        project_id: str | None = None,
        jwt_token: str | None = None,
        url: str | None = None,
    ) -> dict:
        """
        Synchronous method to retrieve the topology of a specific GNS3 project.

        Args:
            tool_input: Input parameters, typically a dict or Pydantic model
                       containing server_url.
            run_manager: Callback manager for tool run.
            project_id: The UUID of the specific GNS3 project to retrieve
                        topology from.
            jwt_token: JWT token for authentication (used by MCP handlers).
            url: GNS3 server URL (used by MCP handlers).

        Returns:
            dict: A dictionary containing the project ID, name, status, nodes,
                  and links, or an error dictionary if an exception occurs
                  or project_id is not provided.
        """

        # Log received input
        logger.info(
            "Received tool_input: %s, project_id: %s", tool_input, project_id
        )

        try:
            # Validate project_id parameter
            if not project_id:
                logger.error("project_id parameter is required.")
                return {
                    "error": "project_id parameter is required. "
                    "Please provide a valid project UUID."
                }

            # Initialize Gns3Connector using factory function
            # jwt_token/url can be passed explicitly (e.g. from MCP handlers)
            # or auto-detected (e.g. from gns3-copilot agent)
            logger.debug("Connecting to GNS3 server...")
            server = get_gns3_connector(jwt_token=jwt_token, url=url)

            if server is None:
                logger.error("Failed to create GNS3 connector")
                return {
                    "error": "Failed to connect to GNS3 server. Please check "
                    "your configuration."
                }

            # Use the provided project_id directly
            logger.info(f"Retrieving topology for project_id: {project_id}")
            project = Project(project_id=project_id, connector=server)
            project.get()  # Load project details

            # Get topology JSON: includes nodes (devices), links, etc.
            topology = {
                "project_id": project.project_id,
                "name": project.name,
                "status": project.status,
                "nodes": self._clean_nodes_ports(
                    copy.deepcopy(project.nodes_inventory())
                ),
                "links": project.links_summary(is_print=False),
            }

            # Log topology result
            logger.info(
                "Topology retrieved: project_id=%s, name=%s, nodes=%d, "
                "links=%d",
                topology.get("project_id"),
                topology.get("name"),
                len(topology.get("nodes", {})),
                len(topology.get("links", [])),
            )
            logger.debug("Topology details: %s", topology)

            return topology

        except Exception as e:
            logger.error("Error retrieving GNS3 topology: %s", str(e))
            return {"error": f"Failed to retrieve topology: {str(e)}"}

    def _clean_nodes_ports(self, data: dict) -> dict:
        """
        Clean and simplify the nodes data structure.

        Simplify each node's ports list to only keep name and short_name
        fields.
        """
        for node in data.values():  # Iterate through R-1, R-2, R-3, R-4
            if "ports" in node and isinstance(node["ports"], list):
                node["ports"] = [
                    {"name": port["name"], "short_name": port["short_name"]}
                    for port in node["ports"]
                ]
        return data


if __name__ == "__main__":
    # Test the tool
    tool = GNS3TopologyTool()

    # Example usage with project_id
    # Replace with an actual project UUID from your GNS3 server
    example_project_id = "0c0fde25-6ead-4413-a283-ea8fd2324291"

    print("Testing GNS3TopologyTool with project_id...")
    result = tool._run(project_id=example_project_id)
    pprint(result)

    # Test without project_id (should return error)
    print("\nTesting without project_id (should return error)...")
    error_result = tool._run()
    pprint(error_result)