opnsense-mcp/src/opnsense_mcp/tools/nat.py

"""NAT rule tools: source NAT, destination NAT, 1:1, NPT."""

from __future__ import annotations

import json
from typing import Any

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.guards import is_dry_run_response, require_writes


def _search_nat(nat_type: str, search_phrase: str, row_count: int) -> str:
    client = get_client()
    try:
        result = client.search_get("firewall", nat_type, "searchRule", search_phrase, row_count)
    except OpnsenseAPIError:
        result = client.search_grid("firewall", nat_type, "searchRule", search_phrase, row_count)
    return json.dumps(result, indent=2, default=str)


def _get_nat(nat_type: str, uuid: str) -> str:
    client = get_client()
    result = client.get("firewall", nat_type, "getRule", uuid)
    return json.dumps(result, indent=2, default=str)


def _add_nat(nat_type: str, action: str, rule: dict[str, Any]) -> str:
    payload = {"rule": rule}
    if settings.dry_run:
        return json.dumps(is_dry_run_response(action, payload), indent=2)
    require_writes(action, payload)
    client = get_client()
    result = client.post("firewall", nat_type, "addRule", data=payload)
    audit_log(action, {"rule": rule, "result": result})
    return json.dumps(result, indent=2, default=str)


def _update_nat(nat_type: str, action: str, uuid: str, rule: dict[str, Any]) -> str:
    payload = {"rule": rule}
    if settings.dry_run:
        return json.dumps(is_dry_run_response(action, {"uuid": uuid, **payload}), indent=2)
    require_writes(action, {"uuid": uuid, **payload})
    client = get_client()
    result = client.post("firewall", nat_type, "setRule", uuid, data=payload)
    audit_log(action, {"uuid": uuid, "rule": rule, "result": result})
    return json.dumps(result, indent=2, default=str)


def _delete_nat(nat_type: str, action: str, uuid: str) -> str:
    if settings.dry_run:
        return json.dumps(is_dry_run_response(action, {"uuid": uuid}), indent=2)
    require_writes(action, {"uuid": uuid})
    client = get_client()
    result = client.post("firewall", nat_type, "delRule", uuid)
    audit_log(action, {"uuid": uuid, "result": result})
    return json.dumps(result, indent=2, default=str)


def _toggle_nat(nat_type: str, action: str, uuid: str, enabled: bool) -> str:
    payload = {"uuid": uuid, "enabled": enabled}
    if settings.dry_run:
        return json.dumps(is_dry_run_response(action, payload), indent=2)
    require_writes(action, payload)
    client = get_client()
    flag = "1" if enabled else "0"
    result = client.post("firewall", nat_type, "toggleRule", uuid, flag)
    audit_log(action, payload | {"result": result})
    return json.dumps(result, indent=2, default=str)


def register(mcp: FastMCP) -> None:
    # ── Source NAT ──────────────────────────────────────────────────────
    @mcp.tool()
    def opnsense_search_source_nat_rules(search_phrase: str = "", row_count: int = 100) -> str:
        """Search outbound / source NAT rules."""
        return _search_nat("source_nat", search_phrase, row_count)

    @mcp.tool()
    def opnsense_get_source_nat_rule(uuid: str) -> str:
        """Get an outbound NAT rule by UUID."""
        return _get_nat("source_nat", uuid)

    @mcp.tool()
    def opnsense_add_source_nat_rule(rule: dict[str, Any]) -> str:
        """Add an outbound / source NAT rule."""
        return _add_nat("source_nat", "add_source_nat_rule", rule)

    @mcp.tool()
    def opnsense_update_source_nat_rule(uuid: str, rule: dict[str, Any]) -> str:
        """Update an outbound NAT rule by UUID."""
        return _update_nat("source_nat", "update_source_nat_rule", uuid, rule)

    @mcp.tool()
    def opnsense_delete_source_nat_rule(uuid: str) -> str:
        """Delete an outbound NAT rule by UUID."""
        return _delete_nat("source_nat", "delete_source_nat_rule", uuid)

    @mcp.tool()
    def opnsense_toggle_source_nat_rule(uuid: str, enabled: bool = True) -> str:
        """Enable or disable an outbound NAT rule."""
        return _toggle_nat("source_nat", "toggle_source_nat_rule", uuid, enabled)

    # ── Destination NAT / Port Forwards ───────────────────────────────────
    @mcp.tool()
    def opnsense_search_d_nat_rules(search_phrase: str = "", row_count: int = 100) -> str:
        """Search destination NAT / port forward rules."""
        return _search_nat("d_nat", search_phrase, row_count)

    @mcp.tool()
    def opnsense_search_port_forwards(search_phrase: str = "", row_count: int = 100) -> str:
        """Search port forward rules (alias for destination NAT search)."""
        return _search_nat("d_nat", search_phrase, row_count)

    @mcp.tool()
    def opnsense_get_d_nat_rule(uuid: str) -> str:
        """Get a destination NAT / port forward rule by UUID."""
        return _get_nat("d_nat", uuid)

    @mcp.tool()
    def opnsense_add_d_nat_rule(rule: dict[str, Any]) -> str:
        """Add a destination NAT / port forward rule."""
        return _add_nat("d_nat", "add_d_nat_rule", rule)

    @mcp.tool()
    def opnsense_update_d_nat_rule(uuid: str, rule: dict[str, Any]) -> str:
        """Update a destination NAT rule by UUID."""
        return _update_nat("d_nat", "update_d_nat_rule", uuid, rule)

    @mcp.tool()
    def opnsense_delete_d_nat_rule(uuid: str) -> str:
        """Delete a destination NAT rule by UUID."""
        return _delete_nat("d_nat", "delete_d_nat_rule", uuid)

    @mcp.tool()
    def opnsense_toggle_d_nat_rule(uuid: str, enabled: bool = True) -> str:
        """Enable or disable a destination NAT rule."""
        return _toggle_nat("d_nat", "toggle_d_nat_rule", uuid, enabled)

    # ── 1:1 NAT ─────────────────────────────────────────────────────────
    @mcp.tool()
    def opnsense_search_one_to_one_rules(search_phrase: str = "", row_count: int = 100) -> str:
        """Search 1:1 NAT rules."""
        return _search_nat("one_to_one", search_phrase, row_count)

    @mcp.tool()
    def opnsense_get_one_to_one_rule(uuid: str) -> str:
        """Get a 1:1 NAT rule by UUID."""
        return _get_nat("one_to_one", uuid)

    @mcp.tool()
    def opnsense_add_one_to_one_rule(rule: dict[str, Any]) -> str:
        """Add a 1:1 NAT rule."""
        return _add_nat("one_to_one", "add_one_to_one_rule", rule)

    @mcp.tool()
    def opnsense_delete_one_to_one_rule(uuid: str) -> str:
        """Delete a 1:1 NAT rule by UUID."""
        return _delete_nat("one_to_one", "delete_one_to_one_rule", uuid)

    # ── NPT (IPv6) ──────────────────────────────────────────────────────
    @mcp.tool()
    def opnsense_search_npt_rules(search_phrase: str = "", row_count: int = 100) -> str:
        """Search NPT (IPv6 prefix translation) rules."""
        return _search_nat("npt", search_phrase, row_count)

    @mcp.tool()
    def opnsense_get_npt_rule(uuid: str) -> str:
        """Get an NPT rule by UUID."""
        return _get_nat("npt", uuid)