Cross-cutting S1 Pydantic models, S3 @write_tool decorator

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

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] = {}

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(
+        result = TestConnectionResult(
-            {
+            status="ok",
-                "status": "ok",
+            url=settings.url,
-                "url": settings.url,
+            firmware=firmware,
-                "firmware": firmware,
+            gateways=gateways,
-                "gateways": gateways,
-            },
-            indent=2,
-            default=str,
         )
+        return result.model_dump_json(indent=2)
     except OpnsenseAPIError as exc:
-        return json.dumps(
+        result = TestConnectionResult(
-            {
+            status="error",
-                "status": "error",
+            url=settings.url,
-                "url": settings.url,
+            firmware={"message": str(exc), "status_code": exc.status_code, "body": exc.body},
-                "message": str(exc),
-                "status_code": exc.status_code,
-                "body": exc.body,
-            },
-            indent=2,
-            default=str,
         )
+        return result.model_dump_json(indent=2)
 
 
 register_all(mcp)

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)
+        result = client.post("firewall", "filter", "addRule", data=_rule_payload(rule))
-        audit_log("add_firewall_rule", {"rule": rule, "result": result})
+        return {"_payload": {"rule": rule, "result": result}, **result}
-        return json.dumps(result, indent=2, default=str)
 
     @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)
+        result = client.post("firewall", "filter", "setRule", uuid, data=_rule_payload(rule))
-        audit_log("update_firewall_rule", {"uuid": uuid, "rule": rule, "result": result})
+        return {"_payload": {"uuid": uuid, "rule": rule, "result": result}, **result}
-        return json.dumps(result, indent=2, default=str)
 
     @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 {"_payload": {"uuid": uuid, "result": result}, **result}
-        return json.dumps(result, indent=2, default=str)
 
     @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 {"_payload": {"uuid": uuid, "enabled": enabled, "result": result}, **result}
-        return json.dumps(result, indent=2, default=str)
 
     @mcp.tool()
     def opnsense_create_firewall_savepoint() -> str:

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,
+        result = HealthCheckResult(
-            "interfaces": interfaces,
+            status="ok",
-            "services": services,
+            timestamp=datetime.now(timezone.utc).isoformat(),
-            "resources": resources,
+            gateways=gateways,
-            "firmware": firmware,
+            resources=resources,
-            "wireguard": wg_status,
+            services=services if isinstance(services, list) else [],
-        }, indent=2, default=str)
+            firmware=firmware,
+            interface_count=len(interfaces) if isinstance(interfaces, list) else 0,
+        )
+        return result.model_dump_json(indent=2)