"""
firewall.py - firewalld manager for Vacuum Wall SSL proxy firewall appliance.

Wraps firewall-cmd CLI via sudo, manages zones, rules, masquerade/NAT,
and port-forwarding. All mutations are --permanent followed by --reload.

A JSON snapshot of all rules is persisted at DATA_DIR/rules.json so the
Flask UI can inspect or restore previous configurations.
"""

import json
import os
import subprocess
from datetime import UTC
from typing import Any

DATA_DIR: str = "/home/wall/vacuum-wall/data/firewall"
RULES_FILE: str = os.path.join(DATA_DIR, "rules.json")


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------


def _run(cmd: list[str], check: bool = True) -> str:
    """Run a command via subprocess and return its stdout.

    Callers must include ``"sudo"`` as the first argument when the
    command requires elevated privileges.

    Raises:
        RuntimeError: When ``check=True`` and the process exits non-zero.
    """
    result = subprocess.run(cmd, capture_output=True, text=True, check=check)
    return result.stdout.strip()


def _reload() -> None:
    """Reload firewalld so permanent changes take effect immediately."""
    _run(["sudo", "firewall-cmd", "--reload"])


def _ensure_data_dir() -> None:
    """Create the data directory tree if it does not exist."""
    os.makedirs(DATA_DIR, exist_ok=True)


# ---------------------------------------------------------------------------
# Read-only queries
# ---------------------------------------------------------------------------


def get_available_zones() -> list[str]:
    """Return the list of all built-in (available) firewalld zone names."""
    output = _run(["sudo", "firewall-cmd", "--get-zones"])
    return output.split()


def get_active_zones() -> dict[str, list[str]]:
    """Return a dict mapping active zone names to their assigned interfaces.

    Example return value::

        {
            "public": ["eth0"],
            "internal": ["eth1"],
        }
    """
    output = _run(["sudo", "firewall-cmd", "--get-active-zones"])
    zones: dict[str, list[str]] = {}
    current_zone: str | None = None
    for raw_line in output.splitlines():
        stripped = raw_line.strip()
        if not stripped:
            continue
        # Indented lines belong to the current zone section.
        if raw_line.startswith(" "):
            current_ifaces = (
                zones[current_zone]
                if current_zone
                else zones.get(list(zones.keys())[-1], [])
            )
            for piece in stripped.split():
                if current_zone and piece not in current_ifaces:
                    current_ifaces.append(piece)
        else:
            current_zone = stripped
            zones[current_zone] = []
    return zones


def get_zone_info(zone: str) -> dict[str, Any]:
    """Return detailed information for *zone*.

    Keys in the returned dict include:
    ``name``, ``target``, ``interfaces``, ``sources``, ``services``,
    ``ports``, ``protocols``, ``forward-ports``, ``masquerade``,
    ``rich-rules``, ``ics``, ``icmp-blocks``, ``module``.
    """
    output = _run(["sudo", "firewall-cmd", f"--zone={zone}", "--list-all"])
    info: dict[str, Any] = {"name": zone}
    for line in output.splitlines():
        line = line.strip()
        if not line or ":" not in line:
            continue
        key, _, value = line.partition(":")
        key = key.strip()
        value = value.strip()

        if not value:
            # Lines like  "interfaces:  " or "masquerade: " when disabled
            if key in ("masquerade", "ics"):
                info[key] = False
            else:
                info[key] = []
        else:
            if key in (
                "interfaces",
                "sources",
                "services",
                "ports",
                "protocols",
                "icmp-blocks",
                "module",
            ):
                info[key] = value.split()
            elif key == "forward-ports":
                info[key] = _parse_forward_ports(value)
            elif key in ("masquerade", "ics"):
                info[key] = value.lower() == "yes"
            elif key == "rich-rules":
                # rich-rules can span multiple lines; we'll parse below.
                info[key] = [value] if value else []
            else:
                info[key] = value

    # rich-rules may already have been set; if not, default to empty.
    info.setdefault("rich-rules", [])
    info.setdefault("interfaces", [])
    info.setdefault("sources", [])
    info.setdefault("services", [])
    info.setdefault("ports", [])
    info.setdefault("protocols", [])
    info.setdefault("forward-ports", [])
    info.setdefault("masquerade", False)
    info.setdefault("ics", False)
    info.setdefault("icmp-blocks", [])
    info.setdefault("module", [])
    info.setdefault("target", "default")
    return info


def get_services() -> list[str]:
    """Return the list of available service names known to firewalld."""
    output = _run(["sudo", "firewall-cmd", "--get-services"])
    return output.split()


def get_icmp_blocks() -> list[str]:
    """Return the list of available ICMP block names."""
    output = _run(["sudo", "firewall-cmd", "--get-icmptypes"])
    return output.split()


def get_interfaces() -> list[str]:
    """Return the list of network interfaces visible via iproute2."""
    output = _run(["ip", "-o", "link", "show"])
    ifaces: list[str] = []
    for line in output.splitlines():
        if line:
            # Format: "NUM: NAME: <FLAGS> ..."
            parts = line.split()
            if len(parts) >= 2:
                name = parts[1].rstrip(":")
                ifaces.append(name)
    return ifaces


def get_rich_rules(zone: str) -> list[str]:
    """Return the rich rules defined for *zone* as a list of raw strings."""
    output = _run(["sudo", "firewall-cmd", f"--zone={zone}", "--list-rich-rules"])
    output = output.strip()
    if not output:
        return []
    rules: list[str] = []
    current: list[str] = []
    for line in output.splitlines():
        raw = line.rstrip()
        if not raw.endswith(";"):
            current.append(raw)
        else:
            current.append(raw)
            rules.append(" ".join(current))
            current = []
    if current:
        rules.append(" ".join(current))
    return rules


# ---------------------------------------------------------------------------
# Zone CRUD
# ---------------------------------------------------------------------------


def create_zone(zone: str, target: str = "default") -> None:
    """Create a new permanent zone in firewalld.

    Args:
        zone: Name of the zone to create.

    Raises:
        RuntimeError: If the zone already exists or creation fails.
    """
    _run(
        [
            "sudo",
            "firewall-cmd",
            f"--zone={zone}",
            f"--set-target={target}",
            "--permanent",
        ]
    )
    _reload()


def delete_zone(zone: str) -> None:
    """Delete an existing zone.

    Raises:
        RuntimeError: If the zone does not exist or the deletion fails.
    """
    _run(["sudo", "firewall-cmd", f"--zone={zone}", "--delete", "--permanent"])
    _reload()


# ---------------------------------------------------------------------------
# Interface assignment
# ---------------------------------------------------------------------------


def set_zone_interfaces(zone: str, interfaces: list[str]) -> None:
    """Assign *interfaces* to *zone*, replacing any existing assignments.

    Existing interfaces on the zone are removed first so only the
    provided list remains.
    """
    # Remove current permanent interfaces for this zone.
    try:
        current = get_zone_info(zone).get("interfaces", [])
    except Exception:
        current = []
    for iface in current:
        _run(
            [
                "sudo",
                "firewall-cmd",
                f"--zone={zone}",
                "--remove-interface=" + iface,
                "--permanent",
            ],
            check=False,
        )

    # Add the desired set.
    for iface in interfaces:
        _run(
            [
                "sudo",
                "firewall-cmd",
                f"--zone={zone}",
                "--add-interface=" + iface,
                "--permanent",
            ]
        )
    _reload()


def add_zone_interface(zone: str, iface: str) -> None:
    """Add a single interface to *zone*."""
    _run(
        [
            "sudo",
            "firewall-cmd",
            f"--zone={zone}",
            "--add-interface=" + iface,
            "--permanent",
        ]
    )
    _reload()


def remove_zone_interface(zone: str, iface: str) -> None:
    """Remove a single interface from *zone*."""
    _run(
        [
            "sudo",
            "firewall-cmd",
            f"--zone={zone}",
            "--remove-interface=" + iface,
            "--permanent",
        ]
    )
    _reload()


# ---------------------------------------------------------------------------
# Service management
# ---------------------------------------------------------------------------


def set_zone_services(zone: str, services: list[str]) -> None:
    """Set services for *zone*, replacing any previously allowed services."""
    # Remove all current services.
    current = get_zone_info(zone).get("services", [])
    for svc in current:
        _run(
            [
                "sudo",
                "firewall-cmd",
                f"--zone={zone}",
                f"--remove-service={svc}",
                "--permanent",
            ],
            check=False,
        )

    for svc in services:
        _run(
            [
                "sudo",
                "firewall-cmd",
                f"--zone={zone}",
                f"--add-service={svc}",
                "--permanent",
            ]
        )
    _reload()


def add_zone_service(zone: str, service: str) -> None:
    """Add a single service to *zone*."""
    _run(
        [
            "sudo",
            "firewall-cmd",
            f"--zone={zone}",
            f"--add-service={service}",
            "--permanent",
        ]
    )
    _reload()


def remove_zone_service(zone: str, service: str) -> None:
    """Remove a single service from *zone*."""
    _run(
        [
            "sudo",
            "firewall-cmd",
            f"--zone={zone}",
            f"--remove-service={service}",
            "--permanent",
        ]
    )
    _reload()


# ---------------------------------------------------------------------------
# Rich rules
# ---------------------------------------------------------------------------


def add_rich_rule(zone: str, rule: str) -> None:
    """Add a rich rule to *zone*.

    The *rule* argument should be a fully-formed rich-rule expression,
    e.g. ``rule family="ipv4" port protocol="tcp" port="443" accept``.
    """
    _run(
        [
            "sudo",
            "firewall-cmd",
            f"--zone={zone}",
            "--add-rich-rule=" + rule,
            "--permanent",
        ]
    )
    _reload()


def remove_rich_rule(zone: str, rule: str) -> None:
    """Remove a rich rule from *zone*.

    The rule string must match exactly what was added.
    """
    _run(
        [
            "sudo",
            "firewall-cmd",
            f"--zone={zone}",
            "--remove-rich-rule=" + rule,
            "--permanent",
        ]
    )
    _reload()


# ---------------------------------------------------------------------------
# Masquerade (NAT)
# ---------------------------------------------------------------------------


def set_masquerade(zone: str, enable: bool) -> None:
    """Enable or disable masquerade (source-NAT) on *zone*."""
    action = "--add-masquerade" if enable else "--remove-masquerade"
    _run(["sudo", "firewall-cmd", f"--zone={zone}", action, "--permanent"])
    _reload()


# ---------------------------------------------------------------------------
# Port forwarding
# ---------------------------------------------------------------------------


def add_forward_port(
    zone: str,
    port: int,
    protocol: str,
    toaddr: str | None = None,
    toport: int | None = None,
) -> None:
    """Add a port forwarding rule to *zone*.

    Forward traffic arriving on ``port/protocol`` to
    ``toaddr:toport`` (or just ``toport`` when *toaddr* is omitted).
    """
    fwd = f"port={port}/proto={protocol}"
    if toaddr and toport:
        fwd += f"/toaddr={toaddr}/toport={toport}"
    elif toport:
        fwd += f"/toport={toport}"
    else:
        fwd += f"/toaddr={toaddr}" if toaddr else ""

    _run(
        [
            "sudo",
            "firewall-cmd",
            f"--zone={zone}",
            f"--add-forward-port={fwd}",
            "--permanent",
        ]
    )
    _reload()


def remove_forward_port(
    zone: str,
    port: int,
    protocol: str,
    toaddr: str | None = None,
    toport: int | None = None,
) -> None:
    """Remove a previously added port-forwarding rule from *zone*.

    All parameters must match the original rule exactly.
    """
    fwd = f"port={port}/proto={protocol}"
    if toaddr and toport:
        fwd += f"/toaddr={toaddr}/toport={toport}"
    elif toport:
        fwd += f"/toport={toport}"
    else:
        fwd += f"/toaddr={toaddr}" if toaddr else ""

    _run(
        [
            "sudo",
            "firewall-cmd",
            f"--zone={zone}",
            f"--remove-forward-port={fwd}",
            "--permanent",
        ]
    )
    _reload()


# ---------------------------------------------------------------------------
# Helpers for parsing forward-port lines
# ---------------------------------------------------------------------------


def _parse_forward_ports(value: str) -> list[str]:
    """Parse the 'forward-ports' line into individual forward-port specifiers.

    Multiple entries are space-separated; each looks like
    ``port=443/proto=tcp/toaddr=192.168.1.5/toport=8080``.
    """
    return value.split() if value else []


# ---------------------------------------------------------------------------
# State snapshot / backup helpers
# ---------------------------------------------------------------------------


def get_state() -> dict[str, Any]:
    """Return the complete current state of firewalld as a Python dict.

    The dict contains all zones with their per-zone configuration, all
    rich rules, masquerade settings, forward-port rules, and the set of
    active interfaces.
    """
    zones: dict[str, dict[str, Any]] = {}
    for name in get_available_zones():
        try:
            zones[name] = get_zone_info(name)
        except Exception:
            continue

    return {
        "active_zones": get_active_zones(),
        "interfaces": get_interfaces(),
        "available_services": get_services(),
        "zones": zones,
        "rich_rules": {name: get_rich_rules(name) for name in zones},
        "timestamp": _now_iso(),
    }


def _now_iso() -> str:
    """Return the current UTC time as an ISO-8601 string."""
    from datetime import datetime

    return datetime.now(UTC).isoformat()


def save_backup() -> str:
    """Capture the full state and write it to RULES_FILE on disk.

    Returns:
        Absolute path to the written file.
    """
    _ensure_data_dir()
    state = get_state()
    with open(RULES_FILE, "w") as fh:
        json.dump(state, fh, indent=2, default=str)
    return RULES_FILE


def load_backup() -> dict[str, Any]:
    """Read the JSON backup file and return the state dict.

    Use :func:`restore_backup` to actually apply the loaded state.

    Raises:
        FileNotFoundError: When no backup file exists at RULES_FILE.
        json.JSONDecodeError: When the file is not valid JSON.

    Returns:
        The loaded state dict.
    """
    with open(RULES_FILE) as fh:
        state: dict[str, Any] = json.load(fh)
    return state


def restore_backup(state: dict[str, Any]) -> None:
    """Apply the zone configuration described in *state*.

    Walks every zone in *state*["zones"] and re-creates services,
    interfaces, forward ports, masquerade, and rich rules.

    This is a *merge*: zones not present in the snapshot are **not**
    touched.
    """
    zones_cfg = state.get("zones", {})
    for zone_name, zinfo in zones_cfg.items():
        # Ensure the zone exists.
        if zone_name not in get_available_zones():
            target = zinfo.get("target", "default")
            create_zone(zone_name, target)

        # Services
        services = zinfo.get("services", [])
        set_zone_services(zone_name, services)

        # Interfaces
        interfaces = zinfo.get("interfaces", [])
        set_zone_interfaces(zone_name, interfaces)

        # Masquerade
        if zinfo.get("masquerade"):
            set_masquerade(zone_name, True)

        # Forward ports  (stored as raw strings in zinfo)
        for fp in zinfo.get("forward-ports", []):
            _run(
                [
                    "sudo",
                    "firewall-cmd",
                    f"--zone={zone_name}",
                    f"--add-forward-port={fp}",
                    "--permanent",
                ],
                check=False,
            )

        # Rich rules
        for rule in zinfo.get("rich-rules", []):
            _run(
                [
                    "sudo",
                    "firewall-cmd",
                    f"--zone={zone_name}",
                    f"--add-rich-rule={rule}",
                    "--permanent",
                ],
                check=False,
            )

    _reload()


__all__ = [
    "DATA_DIR",
    "RULES_FILE",
    "_reload",
    "_run",
    "add_forward_port",
    "add_rich_rule",
    "add_zone_interface",
    "add_zone_service",
    "create_zone",
    "delete_zone",
    "get_active_zones",
    "get_available_zones",
    "get_icmp_blocks",
    "get_interfaces",
    "get_rich_rules",
    "get_services",
    "get_state",
    "get_zone_info",
    "load_backup",
    "remove_forward_port",
    "remove_rich_rule",
    "remove_zone_interface",
    "remove_zone_service",
    "restore_backup",
    "save_backup",
    "set_masquerade",
    "set_zone_interfaces",
    "set_zone_services",
]