FabledSteward/plugins/unifi/client.py

# plugins/unifi/client.py
"""Async UniFi Network controller API client.

Supports UDM/UDM-Pro (firmware 2+) using the /proxy/network/ path prefix
and Bearer-token auth via the /api/auth/login endpoint.

On 401, re-authenticates automatically once before raising.
"""
from __future__ import annotations
import logging

import httpx

logger = logging.getLogger(__name__)


class UnifiClient:
    def __init__(
        self,
        host: str,
        username: str,
        password: str,
        site: str = "default",
        verify_ssl: bool = False,
    ) -> None:
        self._host = host.rstrip("/")
        self._username = username
        self._password = password
        self._site = site
        self._verify_ssl = verify_ssl
        self._http: httpx.AsyncClient | None = None
        self._csrf_token: str = ""

    # ── Lifecycle ─────────────────────────────────────────────────────────────

    async def _ensure_http(self) -> None:
        if self._http is None:
            self._http = httpx.AsyncClient(
                verify=self._verify_ssl,
                timeout=15.0,
                follow_redirects=True,
            )

    async def login(self) -> None:
        await self._ensure_http()
        resp = await self._http.post(
            f"{self._host}/api/auth/login",
            json={"username": self._username, "password": self._password, "rememberMe": False},
        )
        resp.raise_for_status()
        self._csrf_token = resp.headers.get("X-CSRF-Token", "")
        logger.debug("UniFi login OK (csrf=%s…)", self._csrf_token[:8] if self._csrf_token else "none")

    async def close(self) -> None:
        if self._http:
            await self._http.aclose()
            self._http = None

    # ── API helpers ───────────────────────────────────────────────────────────

    def _api_url(self, path: str) -> str:
        return f"{self._host}/proxy/network/api/s/{self._site}{path}"

    def _headers(self) -> dict:
        return {"X-CSRF-Token": self._csrf_token} if self._csrf_token else {}

    async def _get(self, path: str, params: dict | None = None) -> list | dict:
        await self._ensure_http()
        resp = await self._http.get(self._api_url(path), headers=self._headers(), params=params)
        if resp.status_code == 401:
            await self.login()
            resp = await self._http.get(self._api_url(path), headers=self._headers(), params=params)
        resp.raise_for_status()
        data = resp.json()
        return data.get("data", data)

    # ── Core monitoring ───────────────────────────────────────────────────────

    async def get_health(self) -> list[dict]:
        """Subsystem health (wan, wlan, lan, www)."""
        result = await self._get("/stat/health")
        return result if isinstance(result, list) else []

    async def get_active_clients(self) -> list[dict]:
        """Currently connected clients."""
        result = await self._get("/stat/sta")
        return result if isinstance(result, list) else []

    async def get_devices(self) -> list[dict]:
        """Managed network devices (APs, switches, gateways) with full detail."""
        result = await self._get("/stat/device")
        return result if isinstance(result, list) else []

    # ── Events & alarms ───────────────────────────────────────────────────────

    async def get_events(self, limit: int = 200) -> list[dict]:
        """Recent site events (client connects, IDS alerts, device state changes, etc.)."""
        result = await self._get("/stat/event", params={"_limit": limit, "_sort": "-time"})
        return result if isinstance(result, list) else []

    async def get_alarms(self, archived: bool = False) -> list[dict]:
        """Active (or archived) alarms."""
        result = await self._get("/stat/alarm")
        if not isinstance(result, list):
            return []
        return [a for a in result if a.get("archived", False) == archived]

    # ── Traffic & DPI ─────────────────────────────────────────────────────────

    async def get_dpi_stats(self) -> list[dict]:
        """DPI per-client application/category traffic stats (cumulative counters)."""
        result = await self._get("/stat/dpi")
        return result if isinstance(result, list) else []

    async def get_dpi_site_stats(self) -> list[dict]:
        """DPI site-level category totals (not per-client)."""
        result = await self._get("/dpi")
        return result if isinstance(result, list) else []

    async def get_speedtest_status(self) -> dict | None:
        """Most recent WAN speed test result."""
        result = await self._get("/stat/speedtest-status")
        if isinstance(result, list) and result:
            return result[0]
        if isinstance(result, dict):
            return result
        return None

    # ── Network inventory ─────────────────────────────────────────────────────

    async def get_known_clients(self) -> list[dict]:
        """All historically seen clients (not just currently connected)."""
        result = await self._get("/rest/user")
        return result if isinstance(result, list) else []

    async def get_networks(self) -> list[dict]:
        """Network/VLAN configurations."""
        result = await self._get("/rest/networkconf")
        return result if isinstance(result, list) else []

    async def get_wlan_configs(self) -> list[dict]:
        """WLAN (SSID) configurations."""
        result = await self._get("/rest/wlanconf")
        return result if isinstance(result, list) else []

    async def get_port_forwards(self) -> list[dict]:
        """Port forwarding rules."""
        result = await self._get("/rest/portforwarding")
        return result if isinstance(result, list) else []

    async def get_firewall_rules(self) -> list[dict]:
        """Firewall rules."""
        result = await self._get("/rest/firewallrule")
        return result if isinstance(result, list) else []

    async def get_firewall_groups(self) -> list[dict]:
        """Firewall groups (IP sets / port sets referenced by rules)."""
        result = await self._get("/rest/firewallgroup")
        return result if isinstance(result, list) else []

    # ── System ────────────────────────────────────────────────────────────────

    async def get_sysinfo(self) -> dict | None:
        """Controller system information (version, uptime, hostname)."""
        result = await self._get("/stat/sysinfo")
        if isinstance(result, list) and result:
            return result[0]
        if isinstance(result, dict):
            return result
        return None