Agentic-OS/backend/app/agent/playbooks.py

"""Per-device-type connection and read-only diagnostic playbooks.

Each entry maps a DeviceType to:
  - mcp_server: default MCP server that drives the device
  - connect: (tool_name, arg_builder) to establish a session (optional)
  - diagnostics: list of (tool_name, arg_builder) read-only calls to gather state

arg_builder is a callable(device_ctx) -> dict, where device_ctx contains
address, port, username, secret.
"""
from __future__ import annotations

from typing import Callable

from app.agent.asterisk_profiles import ASTERISK_PLAYBOOKS
from app.agent.device_connect import pfsense_connect_args, ssh_connect_args
from app.models.enums import DeviceType

ArgBuilder = Callable[[dict], dict]


def _ip(ctx: dict) -> dict:
    return {"ip": ctx["address"]}


def _pfsense_connect(ctx: dict) -> dict:
    return pfsense_connect_args(ctx)


def _ssh(ctx: dict) -> dict:
    return ssh_connect_args(ctx)


PLAYBOOKS: dict[DeviceType, dict] = {
    DeviceType.pfsense: {
        "mcp_server": "pfsense-mcp",
        "connect": ("pfsense_connect", _pfsense_connect),
        "diagnostics": [
            ("pfsense_get_system_status", lambda c: {}),
            ("pfsense_list_gateways", lambda c: {}),
            ("pfsense_get_gateway_status", lambda c: {}),
            ("pfsense_list_interfaces", lambda c: {}),
        ],
    },
    DeviceType.proxmox: {
        # Connect + diagnostics handled by proxmox_profiles.run_proxmox_diagnostics (API → SSH).
        "mcp_server": "proxmox-mcp",
        "connect": None,
        "diagnostics": [],
    },
    DeviceType.asterisk: {
        # Legacy key — prefer asterisk_geneseasx catalog slot (see ASTERISK_PLAYBOOKS).
        "mcp_server": "ssh-generic-mcp",
        "connect": ("ssh_connect", _ssh),
        "diagnostics": [],
    },
    # GeneseasX / satbox variants are resolved by catalog_key in playbook_for().
    DeviceType.debian: {
        "mcp_server": "ssh-generic-mcp",
        "connect": (
            "ssh_connect",
            lambda c: ssh_connect_args(c),
        ),
        "diagnostics": [
            ("ssh_run", lambda c: {"command": "uptime && uname -a"}),
            ("ssh_run", lambda c: {"command": "df -h"}),
            ("ssh_run", lambda c: {"command": "free -m"}),
            ("ssh_run", lambda c: {"command": "systemctl --failed --no-pager"}),
            ("ssh_run", lambda c: {"command": "journalctl -p err -n 100 --no-pager"}),
        ],
    },
    DeviceType.geneseasx: {
        "mcp_server": "ssh-generic-mcp",
        "connect": (
            "ssh_connect",
            lambda c: ssh_connect_args(c),
        ),
        "diagnostics": [
            ("ssh_run", lambda c: {"command": "uptime && uname -a"}),
            ("ssh_run", lambda c: {"command": "docker ps --format '{{.Names}}: {{.Status}}'"}),
            ("ssh_run", lambda c: {"command": "systemctl --failed --no-pager"}),
        ],
    },
    DeviceType.fortigate: {
        "mcp_server": "fortigate-mcp",
        "connect": (
            "fortigate_connect",
            lambda c: {
                "host": c["address"],
                "api_key": c.get("secret"),
                "username": c.get("username"),
            },
        ),
        "diagnostics": [
            ("fortigate_get_system_status", lambda c: {}),
            ("fortigate_list_interfaces", lambda c: {}),
            ("fortigate_list_policies", lambda c: {}),
            ("fortigate_get_logs", lambda c: {"lines": 100}),
        ],
    },
    DeviceType.fortiswitch: {
        "mcp_server": "fortiswitch-mcp",
        "connect": (
            "fortiswitch_connect",
            lambda c: {
                "host": c["address"],
                "username": c.get("username"),
                "password": c.get("secret"),
            },
        ),
        "diagnostics": [
            ("fortiswitch_get_system_status", lambda c: {}),
            ("fortiswitch_list_ports", lambda c: {}),
            ("fortiswitch_get_port_stats", lambda c: {}),
        ],
    },
}


def playbook_for(
    device_type: DeviceType,
    mcp_override: str | None,
    catalog_key: str | None = None,
) -> dict | None:
    if catalog_key and catalog_key in ASTERISK_PLAYBOOKS:
        pb = ASTERISK_PLAYBOOKS[catalog_key]
    else:
        pb = PLAYBOOKS.get(device_type)
    if pb is None:
        return None
    # Catalog playbooks define the MCP server; ignore stale device-row overrides.
    if mcp_override and catalog_key not in ASTERISK_PLAYBOOKS and device_type != DeviceType.proxmox:
        pb = {**pb, "mcp_server": mcp_override}
    return pb