FabledSteward/plugins/homeassistant/metrics.py

# plugins/homeassistant/metrics.py
"""Home Assistant's metrics, read with the one sign-in Steward already holds, so a
single Home Assistant user covers everything. The scrape stays a Prometheus target:
this lends it a token at each scrape, sends that token only to this Home Assistant,
and stops it when the sign-in has been refused."""
from __future__ import annotations

import asyncio
from datetime import datetime, timezone
from typing import TYPE_CHECKING

from sqlalchemy import or_, select, update

from steward.core.plugin_manager import loaded_plugin_submodule
from steward.core.prometheus import TokenUnavailable

from . import client
from .models import HomeAssistantConnection
from .service import SOURCE

if TYPE_CHECKING:
    from quart import Quart

METRICS_PATH = "/api/prometheus"
NAME = "HomeAssistant"
INTERVAL_SECONDS = 120
TRY_SECONDS = 30
# What a new target leaves out: bookkeeping on every entity, a few thousand series
# that say nothing the readings don't.
DROP = ["homeassistant_state_change*", "homeassistant_last_updated_time_seconds", "homeassistant_entity_available",
        "homeassistant_entity_info", "homeassistant_area_info", "python_*", "process_*"]
REFUSED = "Home Assistant refused Steward's sign-in while reading its metrics, or has banned this address: look in its ip_bans.yaml."


class CannotRead(Exception):
    """Why the metrics can't be read with this sign-in, said to the person."""


def metrics_url(connection) -> str:
    return connection.url.rstrip("/") + METRICS_PATH


# ── Lending the sign-in to a scrape ──────────────────────────────────────────

async def scrape_token(app: "Quart", url: str) -> str:
    """A token for scraping `url`. Raises TokenUnavailable, asking nothing of
    anyone, once the sign-in was refused: another refused request counts toward a
    ban. A target edited to point elsewhere isn't sent this user's token."""
    async with app.db_sessionmaker() as db:
        connection = (await db.execute(select(HomeAssistantConnection))).scalars().first()
    if connection is None:
        raise TokenUnavailable("Home Assistant isn't connected: sign in on its page to read its metrics.")
    if connection.refused:
        raise TokenUnavailable(f"{connection.last_error} Steward has stopped asking until you sign in again "
                               "or choose Try again on the Home Assistant page.")
    if url != metrics_url(connection):
        raise TokenUnavailable(f"This isn't the address of the Home Assistant Steward signed in to, so its sign-in "
                               f"isn't sent here. Set the address to {metrics_url(connection)}, or type a token.")
    try:
        return await client.access_token(connection)
    except (client.Refused, client.Unreachable) as exc:
        raise TokenUnavailable(str(exc)) from None


async def token_refused(app: "Quart") -> None:
    """A scrape's token was refused: stop everything using this sign-in."""
    client.forget_access()
    async with app.db_sessionmaker() as db:
        async with db.begin():
            connection = (await db.execute(select(HomeAssistantConnection))).scalars().first()
            if connection is not None:
                connection.refused, connection.last_error = True, REFUSED


# ── The Metrics section of the page ──────────────────────────────────────────

def _prometheus(sub: str):
    return loaded_plugin_submodule("prometheus", sub)


async def _target(db, connection):
    """The target using this sign-in, or else the one at this address."""
    T = _prometheus("models").PrometheusTarget
    rows = (await db.execute(select(T).where(or_(T.token_from == SOURCE, T.url == metrics_url(connection))))).scalars().all()
    return next((t for t in rows if t.token_from == SOURCE), rows[0] if rows else None)


async def state(db, connection) -> dict | None:
    """What the page says about the metrics. None when the Prometheus plugin is off;
    `id` is None when no target reads them yet."""
    service = _prometheus("service")
    if service is None or _prometheus("models") is None:
        return None
    target = await _target(db, connection)
    if target is None:
        return {"id": None, "name": None, "uses_sign_in": False, "reach": None, "says": None, "kept": 0}
    found = service.target_state(target, datetime.now(timezone.utc))
    return {"id": target.id, "name": target.name, "uses_sign_in": target.token_from == SOURCE,
            "reach": found.reach, "says": found.says, "kept": target.last_kept}


async def try_reading(app: "Quart", connection) -> None:
    """Ask for the metrics once, so turning this on says at once if they aren't there."""
    scraper = _prometheus("scraper")
    if scraper is None:
        raise CannotRead("the Prometheus plugin isn't turned on.")
    if connection.refused:
        raise CannotRead("Home Assistant has refused this sign-in. Choose Try again or sign in again first.")
    try:
        token = await client.access_token(connection)
        await asyncio.wait_for(scraper.scrape(metrics_url(connection), token, connection.verify_tls), TRY_SECONDS)
    except scraper.NoMetricsHere:
        raise CannotRead("Home Assistant isn't serving metrics. Add prometheus: to its configuration.yaml, restart it, "
                         "and come back: Steward doesn't change Home Assistant.") from None
    except scraper.TokenRefused:
        await token_refused(app)
        raise CannotRead(REFUSED) from None
    except TimeoutError:
        raise CannotRead(f"Home Assistant's metrics didn't arrive within {TRY_SECONDS} s.") from None
    except (scraper.ScrapeError, client.Refused, client.Unreachable) as exc:
        raise CannotRead(str(exc)) from None


async def use_sign_in(app: "Quart", connection) -> tuple[str, str, bool]:
    """Make the target, or move the one at this address onto this sign-in, keeping
    its name, filters and history. Returns (id, name, whether it was made)."""
    models, scheduler = _prometheus("models"), _prometheus("scheduler")
    T = models.PrometheusTarget
    async with app.db_sessionmaker() as db:
        async with db.begin():
            target = await _target(db, connection)
            created = target is None
            if created:
                if (await db.execute(select(T.id).where(T.name == NAME))).first():
                    raise CannotRead(f"another Prometheus target is already named {NAME}. Rename it first.")
                target = T(name=NAME, interval_seconds=INTERVAL_SECONDS, drop=list(DROP))
                db.add(target)
            target.url, target.verify_tls = metrics_url(connection), connection.verify_tls
            target.token, target.token_from, target.refused = "", SOURCE, False
            target.last_scraped_at = target.last_error = None
            await db.flush()
            target_id, name = target.id, target.name
    scheduler.RATES.forget(target_id)
    return target_id, name, created


async def follow_sign_in(db, connection) -> None:
    """Signing in at a new address takes the scrape using this sign-in with it."""
    models = _prometheus("models")
    if models is not None:
        T = models.PrometheusTarget
        await db.execute(update(T).where(T.token_from == SOURCE).values(
            url=metrics_url(connection), verify_tls=connection.verify_tls))