Roundtable-plugins/ups/client.py

# plugins/ups/client.py
"""Async NUT (Network UPS Tools) client.

Connects to NUT's TCP port (default 3493) and reads UPS variables.
Read-only: no write commands (test, shutdown via NUT) are implemented.
"""
from __future__ import annotations
import asyncio
import logging

logger = logging.getLogger(__name__)

# NUT status flag meanings
STATUS_FLAGS = {
    "OL":     "Online",
    "OB":     "On Battery",
    "LB":     "Low Battery",
    "CHRG":   "Charging",
    "DISCHRG":"Discharging",
    "BYPASS": "Bypass",
    "CAL":    "Calibrating",
    "FSD":    "Forced Shutdown",
    "RB":     "Replace Battery",
    "OVER":   "Overloaded",
    "TRIM":   "Trimming",
    "BOOST":  "Boosting",
}


class NutError(Exception):
    pass


class NutClient:
    """Minimal async NUT client — connects, reads vars, disconnects."""

    def __init__(
        self,
        host: str = "localhost",
        port: int = 3493,
        username: str = "",
        password: str = "",
        timeout: float = 10.0,
    ) -> None:
        self.host = host
        self.port = port
        self.username = username
        self.password = password
        self.timeout = timeout

    async def get_vars(self, ups_name: str) -> dict[str, str]:
        """Open a connection, authenticate if needed, fetch all UPS variables, close."""
        try:
            reader, writer = await asyncio.wait_for(
                asyncio.open_connection(self.host, self.port),
                timeout=self.timeout,
            )
        except (OSError, asyncio.TimeoutError) as exc:
            raise NutError(f"Cannot connect to NUT at {self.host}:{self.port}: {exc}") from exc

        try:
            async def send(cmd: str) -> None:
                writer.write((cmd + "\n").encode())
                await writer.drain()

            async def readline() -> str:
                line = await asyncio.wait_for(reader.readline(), timeout=self.timeout)
                return line.decode(errors="replace").strip()

            # Optional auth
            if self.username:
                await send(f"USERNAME {self.username}")
                resp = await readline()
                if not resp.startswith("OK"):
                    raise NutError(f"NUT auth USERNAME rejected: {resp}")
            if self.password:
                await send(f"PASSWORD {self.password}")
                resp = await readline()
                if not resp.startswith("OK"):
                    raise NutError(f"NUT auth PASSWORD rejected: {resp}")

            # Request variable list
            await send(f"LIST VAR {ups_name}")

            vars_: dict[str, str] = {}
            while True:
                line = await readline()
                if line.startswith("ERR"):
                    raise NutError(f"NUT error for UPS '{ups_name}': {line}")
                if line.startswith(f"END LIST VAR {ups_name}"):
                    break
                if line.startswith(f"VAR {ups_name} "):
                    # Format: VAR <upsname> <varname> "<value>"
                    rest = line[len(f"VAR {ups_name} "):]
                    if " " in rest:
                        varname, raw_val = rest.split(" ", 1)
                        vars_[varname] = raw_val.strip('"')

            await send("LOGOUT")
            return vars_

        finally:
            writer.close()
            try:
                await writer.wait_closed()
            except Exception:
                pass


def parse_status(vars_: dict[str, str]) -> dict:
    """Extract structured fields from raw NUT variable dict."""
    raw_status = vars_.get("ups.status", "")
    flags = set(raw_status.upper().split())

    def _float(key: str) -> float | None:
        try:
            return float(vars_[key])
        except (KeyError, ValueError):
            return None

    def _int(key: str) -> int | None:
        try:
            return int(float(vars_[key]))
        except (KeyError, ValueError):
            return None

    return {
        "raw_status": raw_status,
        "on_battery": "OB" in flags,
        "low_battery": "LB" in flags,
        "flags": sorted(flags),
        "battery_charge_pct": _float("battery.charge"),
        "battery_runtime_secs": _int("battery.runtime"),
        "battery_voltage": _float("battery.voltage"),
        "load_pct": _float("ups.load"),
        "input_voltage": _float("input.voltage"),
        "output_voltage": _float("output.voltage"),
        "temperature": _float("ups.temperature"),
        "model": vars_.get("ups.model") or vars_.get("device.model"),
        "manufacturer": vars_.get("ups.mfr") or vars_.get("device.mfr"),
    }