FabledSteward/plugins/homeassistant/client.py

# plugins/homeassistant/client.py
"""Talking to Home Assistant, with what its own sign-in gave Steward: a refresh
token kept in the connection, and access tokens made from it that last half an
hour. Repairs and the Supervisor's reports are only offered on its websocket. Home Assistant bans an address after a few refused requests, so nothing is
sent without a token and a refusal stops the asking. Every failure comes back as
one sentence a page can show."""
from __future__ import annotations

import asyncio
import json
import ssl
import time
from contextlib import asynccontextmanager
from urllib.parse import urlencode

import httpx
from websockets.asyncio.client import connect as ws_connect
from websockets.exceptions import WebSocketException

from steward.core.crypto import decrypt_secret

TIMEOUT_SECONDS = 15.0
# Home Assistant answers an action when it has finished, minutes later for a
# backup or an update, and carries on if the request ends. So its answer is only
# waited for this long, and the outcome is read from its entities.
ACTION_READ_SECONDS = 10.0
MAX_BYTES = 16 * 1024 * 1024
# An access token this close to running out isn't used for another request.
EXPIRY_MARGIN_SECONDS = 60
# One read of the system over the websocket, however its answers dribble in.
SYSTEM_READ_SECONDS = 30

_access: dict[str, tuple[str, float]] = {}      # refresh token -> access token, when to stop using it


class Unreachable(Exception):
    """Home Assistant didn't answer usefully; it may be restarting."""


class Refused(Exception):
    """Home Assistant refused Steward's sign-in, or has banned this address."""


class Rejected(Exception):
    """Home Assistant understood an action and wouldn't do it."""


def _refusal(status: int) -> Refused | None:
    if status == 401:
        _access.clear()
        return Refused("Home Assistant refused Steward's sign-in. Sign in again, as a user who may do this.")
    if status == 403:
        return Refused("Home Assistant answered 403: it may have banned Steward's address "
                       "after refused requests. Look in its ip_bans.yaml.")
    return None


async def _request(method: str, url: str, path: str, verify_tls: bool, transport, *, token: str | None = None,
                   read_timeout: float = TIMEOUT_SECONDS, **kwargs) -> tuple[int, bytes]:
    headers = {"Authorization": f"Bearer {token}"} if token else {}
    timeout = httpx.Timeout(TIMEOUT_SECONDS, read=read_timeout)
    try:
        # A slow drip of bytes can keep each read inside its own timeout for ever.
        async with asyncio.timeout(TIMEOUT_SECONDS + read_timeout):
            async with httpx.AsyncClient(timeout=timeout, verify=verify_tls, transport=transport,
                                         follow_redirects=False) as client:
                async with client.stream(method, url.rstrip("/") + path, headers=headers, **kwargs) as response:
                    body = bytearray()
                    async for chunk in response.aiter_bytes():
                        body.extend(chunk)
                        if len(body) > MAX_BYTES:
                            raise Unreachable("Home Assistant's answer was too large to read.")
                    return response.status_code, bytes(body)
    except TimeoutError:
        raise httpx.ReadTimeout("no complete answer in time") from None


def _unreachable(exc: httpx.HTTPError) -> Unreachable:
    if isinstance(exc, httpx.TimeoutException):
        return Unreachable(f"Home Assistant didn't answer within {TIMEOUT_SECONDS:.0f} s.")
    return Unreachable(f"Couldn't reach Home Assistant: {exc}")


# ── Signing in ───────────────────────────────────────────────────────────────

def sign_in_url(url: str, client_id: str, redirect_uri: str, state: str) -> str:
    """Home Assistant's own sign-in page, which sends the browser back to
    `redirect_uri`. It accepts `client_id` as it is because the two share an
    address, so it needn't be able to reach Steward."""
    return f"{url.rstrip('/')}/auth/authorize?" + urlencode(
        {"response_type": "code", "client_id": client_id, "redirect_uri": redirect_uri, "state": state})


async def _tokens(url: str, form: dict, verify_tls: bool, transport) -> dict:
    try:
        status, body = await _request("POST", url, "/auth/token", verify_tls, transport, data=form)
    except httpx.HTTPError as exc:
        raise _unreachable(exc) from None
    if status in (400, 401, 403):
        raise Refused("Home Assistant no longer accepts Steward's sign-in. Sign in again.")
    if status != 200:
        raise Unreachable(f"Home Assistant answered {status}.")
    try:
        found = json.loads(body)
        found["access_token"], int(found["expires_in"])
    except (ValueError, KeyError, TypeError):
        raise Unreachable("Home Assistant's answer wasn't a sign-in.") from None
    return found


def _remember(refresh_token: str, found: dict) -> str:
    until = time.monotonic() + int(found["expires_in"]) - EXPIRY_MARGIN_SECONDS
    _access[refresh_token] = (found["access_token"], until)
    return found["access_token"]


async def redeem_code(url: str, client_id: str, code: str, verify_tls: bool, transport=None) -> tuple[str, str]:
    """(refresh token, access token) for the code the sign-in page sent back."""
    found = await _tokens(url, {"grant_type": "authorization_code", "client_id": client_id, "code": code},
                          verify_tls, transport)
    refresh_token = found.get("refresh_token")
    if not refresh_token:
        raise Unreachable("Home Assistant's answer wasn't a sign-in.")
    return refresh_token, _remember(refresh_token, found)


async def access_token(connection, transport=None) -> str:
    refresh_token = decrypt_secret(connection.refresh_token, context="Home Assistant sign-in")
    token, until = _access.get(refresh_token, (None, 0.0))
    if token and time.monotonic() < until:
        return token
    found = await _tokens(connection.url, {"grant_type": "refresh_token", "client_id": connection.client_id,
                                           "refresh_token": refresh_token}, connection.verify_tls, transport)
    return _remember(refresh_token, found)


async def is_admin(url: str, token: str, verify_tls: bool, transport=None) -> bool:
    """Whether the signed-in user is an administrator, which its API only says by
    refusing someone who isn't. So this is asked once, right after signing in, and
    a refusal here counts once toward a ban."""
    try:
        status, _ = await _request("POST", url, "/api/template", verify_tls, transport, token=token,
                                   json={"template": "1"})
    except httpx.HTTPError as exc:
        raise _unreachable(exc) from None
    if status == 401:
        return False
    if status != 200:
        raise (_refusal(status) or Unreachable(f"Home Assistant answered {status}."))
    return True


async def revoke(connection, transport=None) -> None:
    """End this sign-in on Home Assistant. It answers 200 whether or not it knew it."""
    refresh_token = decrypt_secret(connection.refresh_token, context="Home Assistant sign-in")
    _access.pop(refresh_token, None)
    try:
        status, _ = await _request("POST", connection.url, "/auth/revoke", connection.verify_tls, transport,
                                   data={"token": refresh_token})
    except httpx.HTTPError as exc:
        raise _unreachable(exc) from None
    if status != 200:
        raise Unreachable(f"Home Assistant answered {status}.")


# ── Reading and acting ───────────────────────────────────────────────────────

async def _get(url: str, path: str, token: str, verify_tls: bool, transport) -> list | dict | None:
    """The parsed JSON of a GET, or None for a 404."""
    try:
        status, body = await _request("GET", url, path, verify_tls, transport, token=token)
    except httpx.HTTPError as exc:
        raise _unreachable(exc) from None
    if refusal := _refusal(status):
        raise refusal
    if status == 404:
        return None
    if status != 200:
        raise Unreachable(f"Home Assistant answered {status}.")
    try:
        return json.loads(body)
    except ValueError:
        raise Unreachable("Home Assistant's answer wasn't JSON.") from None


async def read_states(url: str, token: str, verify_tls: bool, transport=None) -> list[dict]:
    states = await _get(url, "/api/states", token, verify_tls, transport)
    if not isinstance(states, list):
        raise Unreachable("Home Assistant's answer wasn't a list of entities.")
    return states


async def read_state(url: str, entity_id: str, token: str, verify_tls: bool, transport=None) -> dict | None:
    return await _get(url, f"/api/states/{entity_id}", token, verify_tls, transport)


async def call_action(url: str, domain: str, action: str, data: dict, token: str, verify_tls: bool,
                      transport=None) -> None:
    """Ask Home Assistant to do something. Returning means it got the request,
    not that it has finished: a request cut off once sent has been taken up, and
    an update that restarts Home Assistant cuts its own request off."""
    try:
        status, body = await _request("POST", url, f"/api/services/{domain}/{action}", verify_tls, transport,
                                      token=token, read_timeout=ACTION_READ_SECONDS, json=data)
    except (httpx.ReadTimeout, httpx.ReadError, httpx.RemoteProtocolError):
        return
    except httpx.HTTPError as exc:
        raise _unreachable(exc) from None
    if refusal := _refusal(status):
        raise refusal
    if status != 200:
        try:
            said = json.loads(body).get("message")
        except (ValueError, AttributeError):
            said = None
        raise Rejected(said or f"Home Assistant answered {status}.")


# ── The system, over the websocket ───────────────────────────────────────────

class _Socket:
    """Home Assistant's websocket after sign-in: one question at a time, each with
    an id, and None for one it won't answer (no Supervisor, or not an administrator)."""

    def __init__(self, ws):
        self.ws, self.asked = ws, 0

    async def sign_in(self, token: str) -> None:
        await self.ws.recv()
        await self.ws.send(json.dumps({"type": "auth", "access_token": token}))
        if json.loads(await self.ws.recv()).get("type") != "auth_ok":
            raise Refused("Home Assistant refused Steward's sign-in. Sign in again, as a user who may do this.")

    async def ask(self, **question):
        self.asked += 1
        await self.ws.send(json.dumps({"id": self.asked, **question}))
        while True:
            answer = json.loads(await self.ws.recv())
            if answer.get("id") == self.asked:
                return answer.get("result") if answer.get("success") else None

    async def supervisor(self, endpoint: str):
        return await self.ask(type="supervisor/api", endpoint=endpoint, method="get")


@asynccontextmanager
async def _signed_in(url: str, token: str, verify_tls: bool, connect):
    """A signed-in socket. `connect` stands in for the websocket in tests."""
    target = "ws" + url.rstrip("/").removeprefix("http") + "/api/websocket"
    options: dict = {"max_size": MAX_BYTES}
    if target.startswith("wss:") and not verify_tls:
        options["ssl"] = ssl.create_default_context()
        options["ssl"].check_hostname, options["ssl"].verify_mode = False, ssl.CERT_NONE
    try:
        async with asyncio.timeout(SYSTEM_READ_SECONDS):
            async with (connect or ws_connect)(target, **options) as ws:
                socket = _Socket(ws)
                await socket.sign_in(token)
                yield socket
    except TimeoutError:
        raise Unreachable(f"Home Assistant's system report didn't arrive within {SYSTEM_READ_SECONDS} s.") from None
    except (OSError, ValueError, KeyError, WebSocketException) as exc:
        raise Unreachable(f"Couldn't read Home Assistant's system report: {exc}") from None


async def read_system(url: str, token: str, verify_tls: bool, connect=None) -> dict:
    """What service.read_health needs, as Home Assistant answered it."""
    async with _signed_in(url, token, verify_tls, connect) as socket:
        found = {"config": await socket.ask(type="get_config"),
                 "issues": (await socket.ask(type="repairs/list_issues") or {}).get("issues"),
                 "resolution": await socket.supervisor("/resolution/info"),
                 "host": await socket.supervisor("/host/info"), "titles": {}, "boots": {}}
        domains = sorted({i["domain"] for i in found["issues"] or []})
        if domains:
            titles = await socket.ask(type="frontend/get_translations", language="en",
                                      category="issues", integration=domains)
            found["titles"] = (titles or {}).get("resources") or {}
        found["addons"] = (await socket.supervisor("/addons") or {}).get("addons")
        # Whether one should be running is only on its own page, so a running one isn't asked.
        for addon in found["addons"] or []:
            if addon.get("state") != "started":
                info = await socket.supervisor(f"/addons/{addon['slug']}/info") or {}
                found["boots"][addon["slug"]] = info.get("boot")
        return found


async def read_os(url: str, token: str, verify_tls: bool, connect=None) -> dict | None:
    """The Supervisor's report on the OS: `version`, and `version_pending` while a
    written update waits for its reboot. None if it isn't answering yet, as just
    after a boot."""
    async with _signed_in(url, token, verify_tls, connect) as socket:
        return await socket.supervisor("/os/info")