From c6f9a20e59b06f6b9e7dc6312912e4d87b0fbfa5 Mon Sep 17 00:00:00 2001 From: Hermes CT107 Date: Sat, 4 Jul 2026 20:51:11 +0000 Subject: [PATCH] Cross-cutting S1 Pydantic models, S3 @write_tool decorator --- src/opnsense_mcp/decorators.py | 61 ++++++++++++++++++++++++++++++ src/opnsense_mcp/schemas.py | 24 ++++++++++++ src/opnsense_mcp/server.py | 31 ++++++--------- src/opnsense_mcp/tools/firewall.py | 47 ++++++++--------------- src/opnsense_mcp/tools/system.py | 22 +++++++---- 5 files changed, 126 insertions(+), 59 deletions(-) create mode 100644 src/opnsense_mcp/decorators.py create mode 100644 src/opnsense_mcp/schemas.py diff --git a/src/opnsense_mcp/decorators.py b/src/opnsense_mcp/decorators.py new file mode 100644 index 0000000..efc7f42 --- /dev/null +++ b/src/opnsense_mcp/decorators.py @@ -0,0 +1,61 @@ +"""Tool decorators for boilerplate reduction in write operations.""" + +from __future__ import annotations + +import functools +import json +from typing import Any, Callable + +from opnsense_mcp.audit import audit_log +from opnsense_mcp.config import settings +from opnsense_mcp.guards import is_dry_run_response, require_writes + + +def write_tool(action: str) -> Callable: + """Decorator that handles dry_run, writes guard, audit log, and JSON encoding. + + The wrapped function should return a dict. If it includes a ``_payload`` key, + that value is used for the audit log entry; otherwise the whole result dict + is used. The ``_payload`` key is stripped from the final JSON output. + + Usage:: + + @write_tool("add_firewall_rule") + def opnsense_add_firewall_rule(rule: dict[str, Any]) -> dict[str, Any]: + client = get_client() + result = client.post(...) + return {"_payload": {"rule": rule, "result": result}, **result} + """ + + def decorator(func: Callable) -> Callable: + @functools.wraps(func) + def wrapper(*args: Any, **kwargs: Any) -> str: + # Build payload from keyword arguments for dry_run / require_writes + payload: dict[str, Any] = dict(kwargs) + + if settings.dry_run: + return json.dumps(is_dry_run_response(action, payload), indent=2) + + require_writes(action, payload) + + result = func(*args, **kwargs) + + # Extract audit payload from result + if isinstance(result, dict): + audit_payload = dict(result.get("_payload", {})) + if not audit_payload: + audit_payload = dict(result) + else: + audit_payload = {"result": result} + + audit_log(action, audit_payload) + + # Strip _payload before returning + if isinstance(result, dict): + result = {k: v for k, v in result.items() if k != "_payload"} + + return json.dumps(result, indent=2, default=str) + + return wrapper + + return decorator diff --git a/src/opnsense_mcp/schemas.py b/src/opnsense_mcp/schemas.py new file mode 100644 index 0000000..8be6f23 --- /dev/null +++ b/src/opnsense_mcp/schemas.py @@ -0,0 +1,24 @@ +"""Pydantic response models for OPNsense MCP tools.""" + +from __future__ import annotations + +from typing import Any + +from pydantic import BaseModel + + +class HealthCheckResult(BaseModel): + status: str + timestamp: str + gateways: dict[str, Any] = {} + resources: dict[str, Any] = {} + services: list[dict[str, Any]] = [] + firmware: dict[str, Any] = {} + interface_count: int = 0 + + +class TestConnectionResult(BaseModel): + status: str + url: str + firmware: dict[str, Any] = {} + gateways: dict[str, Any] = {} diff --git a/src/opnsense_mcp/server.py b/src/opnsense_mcp/server.py index 59e6fea..b9aefee 100644 --- a/src/opnsense_mcp/server.py +++ b/src/opnsense_mcp/server.py @@ -8,6 +8,7 @@ from mcp.server.fastmcp import FastMCP from opnsense_mcp import __version__ from opnsense_mcp.config import settings +from opnsense_mcp.schemas import TestConnectionResult from opnsense_mcp.tools import register_all mcp = FastMCP( @@ -48,28 +49,20 @@ def opnsense_test_connection() -> str: try: firmware = client.get("core", "firmware", "status") gateways = client.get("routes", "gateway", "status") - return json.dumps( - { - "status": "ok", - "url": settings.url, - "firmware": firmware, - "gateways": gateways, - }, - indent=2, - default=str, + result = TestConnectionResult( + status="ok", + url=settings.url, + firmware=firmware, + gateways=gateways, ) + return result.model_dump_json(indent=2) except OpnsenseAPIError as exc: - return json.dumps( - { - "status": "error", - "url": settings.url, - "message": str(exc), - "status_code": exc.status_code, - "body": exc.body, - }, - indent=2, - default=str, + result = TestConnectionResult( + status="error", + url=settings.url, + firmware={"message": str(exc), "status_code": exc.status_code, "body": exc.body}, ) + return result.model_dump_json(indent=2) register_all(mcp) diff --git a/src/opnsense_mcp/tools/firewall.py b/src/opnsense_mcp/tools/firewall.py index dbf88be..d16b4be 100644 --- a/src/opnsense_mcp/tools/firewall.py +++ b/src/opnsense_mcp/tools/firewall.py @@ -10,6 +10,7 @@ from mcp.server.fastmcp import FastMCP from opnsense_mcp.audit import audit_log from opnsense_mcp.client import OpnsenseAPIError, get_client from opnsense_mcp.config import settings +from opnsense_mcp.decorators import write_tool from opnsense_mcp.guards import ( GuardError, is_dry_run_response, @@ -66,57 +67,39 @@ def register(mcp: FastMCP) -> None: return json.dumps(result, indent=2, default=str) @mcp.tool() - def opnsense_add_firewall_rule(rule: dict[str, Any]) -> str: + @write_tool("add_firewall_rule") + def opnsense_add_firewall_rule(rule: dict[str, Any]) -> dict[str, Any]: """Add a new firewall filter rule. Rule dict uses OPNsense field names (description, interface, source_net, destination_net, protocol, action, etc.). Descriptions must start with the configured prefix (default 'mcp:').""" validate_firewall_rule(rule, is_new=True) - payload = _rule_payload(rule) - if settings.dry_run: - return json.dumps(is_dry_run_response("add_firewall_rule", payload), indent=2) - require_writes("add_firewall_rule", payload) client = get_client() - result = client.post("firewall", "filter", "addRule", data=payload) - audit_log("add_firewall_rule", {"rule": rule, "result": result}) - return json.dumps(result, indent=2, default=str) + result = client.post("firewall", "filter", "addRule", data=_rule_payload(rule)) + return {"_payload": {"rule": rule, "result": result}, **result} @mcp.tool() - def opnsense_update_firewall_rule(uuid: str, rule: dict[str, Any]) -> str: + @write_tool("update_firewall_rule") + def opnsense_update_firewall_rule(uuid: str, rule: dict[str, Any]) -> dict[str, Any]: """Update an existing firewall rule by UUID.""" validate_firewall_rule(rule, is_new=False) - payload = _rule_payload(rule) - if settings.dry_run: - return json.dumps( - is_dry_run_response("update_firewall_rule", {"uuid": uuid, **payload}), indent=2 - ) - require_writes("update_firewall_rule", {"uuid": uuid, **payload}) client = get_client() - result = client.post("firewall", "filter", "setRule", uuid, data=payload) - audit_log("update_firewall_rule", {"uuid": uuid, "rule": rule, "result": result}) - return json.dumps(result, indent=2, default=str) + result = client.post("firewall", "filter", "setRule", uuid, data=_rule_payload(rule)) + return {"_payload": {"uuid": uuid, "rule": rule, "result": result}, **result} @mcp.tool() - def opnsense_delete_firewall_rule(uuid: str) -> str: + @write_tool("delete_firewall_rule") + def opnsense_delete_firewall_rule(uuid: str) -> dict[str, Any]: """Delete a firewall rule by UUID.""" - payload = {"uuid": uuid} - if settings.dry_run: - return json.dumps(is_dry_run_response("delete_firewall_rule", payload), indent=2) - require_writes("delete_firewall_rule", payload) client = get_client() result = client.post("firewall", "filter", "delRule", uuid) - audit_log("delete_firewall_rule", {"uuid": uuid, "result": result}) - return json.dumps(result, indent=2, default=str) + return {"_payload": {"uuid": uuid, "result": result}, **result} @mcp.tool() - def opnsense_toggle_firewall_rule(uuid: str, enabled: bool = True) -> str: + @write_tool("toggle_firewall_rule") + def opnsense_toggle_firewall_rule(uuid: str, enabled: bool = True) -> dict[str, Any]: """Enable or disable a firewall rule. enabled=true to enable, false to disable.""" - payload = {"uuid": uuid, "enabled": enabled} - if settings.dry_run: - return json.dumps(is_dry_run_response("toggle_firewall_rule", payload), indent=2) - require_writes("toggle_firewall_rule", payload) client = get_client() flag = "1" if enabled else "0" result = client.post("firewall", "filter", "toggleRule", uuid, flag) - audit_log("toggle_firewall_rule", payload | {"result": result}) - return json.dumps(result, indent=2, default=str) + return {"_payload": {"uuid": uuid, "enabled": enabled, "result": result}, **result} @mcp.tool() def opnsense_create_firewall_savepoint() -> str: diff --git a/src/opnsense_mcp/tools/system.py b/src/opnsense_mcp/tools/system.py index 0334b3a..5b5741f 100644 --- a/src/opnsense_mcp/tools/system.py +++ b/src/opnsense_mcp/tools/system.py @@ -3,11 +3,13 @@ from __future__ import annotations import json +from datetime import datetime, timezone from typing import Any from mcp.server.fastmcp import FastMCP from opnsense_mcp.client import get_client +from opnsense_mcp.schemas import HealthCheckResult def register(mcp: FastMCP) -> None: @@ -67,6 +69,7 @@ def register(mcp: FastMCP) -> None: resources = client.get("diagnostics", "system", "systemResources") memory = client.get("diagnostics", "system", "memory") return json.dumps({"resources": resources, "memory": memory}, indent=2, default=str) + @mcp.tool() def opnsense_health_check() -> str: """Comprehensive health snapshot: gateways, interfaces, services, system resources, firmware, WireGuard status.""" @@ -95,11 +98,14 @@ def register(mcp: FastMCP) -> None: wg_status = client.get("wireguard", "service", "show") except Exception: wg_status = {"error": "WireGuard check failed or not installed"} - return json.dumps({ - "gateways": gateways, - "interfaces": interfaces, - "services": services, - "resources": resources, - "firmware": firmware, - "wireguard": wg_status, - }, indent=2, default=str) + + result = HealthCheckResult( + status="ok", + timestamp=datetime.now(timezone.utc).isoformat(), + gateways=gateways, + resources=resources, + services=services if isinstance(services, list) else [], + firmware=firmware, + interface_count=len(interfaces) if isinstance(interfaces, list) else 0, + ) + return result.model_dump_json(indent=2)