homelab-codex-ws/scripts/ha/lib/ha_ws.py

#!/usr/bin/env python3
"""Minimal synchronous client for HA's `/api/websocket` API.

Used by the "api" adapter (see DESIGN.md, "Deploy path: adapter per
instance") to reach dashboards, area/entity registries, and input_* helpers
— none of these are exposed over REST. Read-only by construction: the only
commands this module ever sends are `auth` and whatever `*_list`/
`lovelace/config` command the caller passes to `HaWsClient.command()` — it
has no method that issues a create/update/delete command.

Depends on the `websocket-client` PyPI package (import name `websocket`),
which is not in the stdlib. See kb/runbooks/home-assistant-deploy.md,
"WebSocket dependency", for the install command. Importing this module
raises ImportError with an actionable message if it's missing, instead of
the default traceback, so the api adapter can surface one clear line
instead of a stack trace.
"""
import itertools
import json

try:
    import websocket as _ws
except ImportError as exc:  # pragma: no cover - exercised via import.sh error path
    raise ImportError(
        "the 'websocket-client' package is required for the api adapter's "
        "dashboard/registry/helper export (HA's WebSocket API) — install it "
        "with 'sudo apt install python3-websocket' or 'pip install --user "
        "websocket-client' (see kb/runbooks/home-assistant-deploy.md, "
        "'WebSocket dependency')"
    ) from exc


class HaWsError(RuntimeError):
    """Raised for handshake/auth failures or a command that reports failure."""


class HaWsClient:
    """One `/api/websocket` connection: authenticate once, issue N read-only commands."""

    def __init__(self, base_url, token, timeout=30):
        ws_url = base_url.rstrip("/")
        if ws_url.startswith("https://"):
            ws_url = "wss://" + ws_url[len("https://"):]
        elif ws_url.startswith("http://"):
            ws_url = "ws://" + ws_url[len("http://"):]
        self._url = ws_url + "/api/websocket"
        self._token = token
        self._timeout = timeout
        self._next_id = itertools.count(1)
        self._conn = None

    def __enter__(self):
        self.connect()
        return self

    def __exit__(self, *exc_info):
        self.close()

    def connect(self):
        self._conn = _ws.create_connection(self._url, timeout=self._timeout)
        hello = self._recv()
        if hello.get("type") != "auth_required":
            self.close()
            raise HaWsError(f"unexpected handshake message from {self._url}: {hello!r}")

        self._conn.send(json.dumps({"type": "auth", "access_token": self._token}))
        auth_result = self._recv()
        if auth_result.get("type") != "auth_ok":
            self.close()
            raise HaWsError(
                "websocket auth rejected — check the deploy_agent token "
                f"(server said: {auth_result!r})"
            )

    def close(self):
        if self._conn is not None:
            self._conn.close()
            self._conn = None

    def _recv(self):
        return json.loads(self._conn.recv())

    def command(self, command_type, **kwargs):
        """Send one read-only WS command, return its `result` payload.

        Raises HaWsError if the server reports failure or sends something
        other than a matching `result` message.
        """
        msg_id = next(self._next_id)
        payload = {"id": msg_id, "type": command_type}
        payload.update(kwargs)
        self._conn.send(json.dumps(payload))
        while True:
            msg = self._recv()
            if msg.get("id") != msg_id:
                # Not our response (e.g. an unrelated event) — keep waiting.
                continue
            if msg.get("type") == "result":
                if not msg.get("success", False):
                    raise HaWsError(f"{command_type} failed: {msg.get('error')}")
                return msg.get("result")
            raise HaWsError(f"unexpected response to {command_type}: {msg!r}")