FabledSteward/plugins/homeassistant/updates.py

# plugins/homeassistant/updates.py
"""Installing one update: take a backup, confirm it landed, install, and follow
it to its end. Its row holds where it has got to, so Steward restarting, or Home
Assistant restarting as an update does, only delays the next step."""
from __future__ import annotations

import logging
from datetime import datetime, timezone

from sqlalchemy import select
from sqlalchemy.exc import IntegrityError

from steward.core.crypto import decrypt_secret

from . import client
from .models import HomeAssistantConnection, HomeAssistantUpdate
from .service import BACKUP_EVENT, LAST_BACKUP, automation_entity_id, is_newer, parse_moment
from .timeline import update_entry

logger = logging.getLogger(__name__)

UNDER_WAY = ("backing_up", "installing")
BACKUP_DEADLINE_S = 30 * 60
# How long the automation has to show it ran before the webhook is called wrong.
TRIGGER_GRACE_S = 90
# An OS update reboots the machine, and a Pi is slow to come back.
INSTALL_DEADLINE_S = {"addon": 20 * 60, "core": 30 * 60, "os": 45 * 60}


class CannotStart(Exception):
    """Why an update wasn't started, said to the person."""


def _secrets(connection) -> tuple[str, str]:
    return (decrypt_secret(connection.token, context="Home Assistant token"),
            decrypt_secret(connection.webhook_id, context="Home Assistant webhook id"))


async def _mark_refused(app, connection_id: str, said: str) -> None:
    async with app.db_sessionmaker() as session:
        async with session.begin():
            row = await session.get(HomeAssistantConnection, connection_id)
            if row is not None:
                row.refused, row.last_error = True, said


async def under_way(session) -> HomeAssistantUpdate | None:
    return (await session.execute(select(HomeAssistantUpdate).where(
        HomeAssistantUpdate.state.in_(UNDER_WAY)))).scalars().first()


async def start_update(app, entity_id: str, username: str, transport=None) -> HomeAssistantUpdate:
    """Start installing `entity_id`, an explicit action that may contact Home
    Assistant. Raises CannotStart, with nothing started, unless the update is
    still waiting and the automation is there to install it."""
    async with app.db_sessionmaker() as session:
        connection = (await session.execute(select(HomeAssistantConnection))).scalars().first()
        if connection is None:
            raise CannotStart("Home Assistant isn't connected")
        if connection.refused:
            raise CannotStart("Home Assistant refused Steward's token")
        if await under_way(session) is not None:
            raise CannotStart("another update is under way")
        if connection.automation != "on":
            raise CannotStart("the Steward maintenance automation isn't on in Home Assistant")
    url, verify = connection.url, connection.verify_tls
    token, webhook_id = _secrets(connection)
    known = next((u for u in connection.updates or [] if u["entity_id"] == entity_id), None)
    if known is None or not known["kind"]:
        raise CannotStart("that isn't a Core, OS or add-on update")
    try:
        current = await client.read_state(url, entity_id, token, verify, transport)
        backup = await client.read_state(url, LAST_BACKUP, token, verify, transport)
        event = await client.read_state(url, BACKUP_EVENT, token, verify, transport)
        automation = await client.read_state(url, automation_entity_id(connection), token, verify, transport)
    except client.Refused as exc:
        await _mark_refused(app, connection.id, str(exc))
        raise CannotStart(str(exc).rstrip(".")) from None
    except client.Unreachable as exc:
        raise CannotStart(str(exc).rstrip(".")) from None
    if backup is None:
        raise CannotStart("Home Assistant has no automatic backups to confirm one with")
    if current is None or current.get("state") != "on":
        raise CannotStart("that update is no longer waiting")
    attrs = current.get("attributes") or {}
    run = HomeAssistantUpdate(
        entity_id=entity_id, title=known["title"], kind=known["kind"],
        from_version=attrs.get("installed_version"), to_version=attrs.get("latest_version"),
        started_by=username, backup_before=backup.get("state"),
        backup_event_before=(event or {}).get("state"),
        triggered_before=((automation or {}).get("attributes") or {}).get("last_triggered"))
    async with app.db_sessionmaker() as session:
        try:
            async with session.begin():
                session.add(run)
                await session.flush()
                run_id = run.id
        except IntegrityError:
            raise CannotStart("another update is under way") from None
    try:
        await client.send_webhook(url, webhook_id, {"action": "backup"}, verify, transport)
    except client.Unreachable as exc:
        await _finish(app, run_id, "failed", f"The backup couldn't be started: {exc} Nothing was installed.")
        raise CannotStart(str(exc).rstrip(".")) from None
    async with app.db_sessionmaker() as session:
        return await session.get(HomeAssistantUpdate, run_id)


async def _finish(app, run_id: str, state: str, message: str | None) -> None:
    async with app.db_sessionmaker() as session:
        async with session.begin():
            run = await session.get(HomeAssistantUpdate, run_id)
            run.state, run.message, run.finished_at = state, message, datetime.now(timezone.utc)
            # The page should show the new versions now, not at the next poll.
            connection = (await session.execute(select(HomeAssistantConnection))).scalars().first()
            if connection is not None:
                connection.last_polled_at = None
            try:
                async with session.begin_nested():
                    await update_entry(session, run)
            except Exception:
                logger.exception("Home Assistant update %s: its timeline entry could not be written", run_id)


async def advance(app, transport=None) -> None:
    """Move the update under way one step on, if there is one."""
    async with app.db_sessionmaker() as session:
        run = await under_way(session)
        connection = (await session.execute(select(HomeAssistantConnection))).scalars().first()
    if run is None:
        return
    if connection is None:
        await _finish(app, run.id, "failed", "The Home Assistant connection was removed.")
        return
    try:
        if run.state == "backing_up":
            await _advance_backup(app, run, connection, transport)
        else:
            await _advance_install(app, run, connection, transport)
    except client.Refused as exc:
        await _mark_refused(app, connection.id, str(exc))
        await _finish(app, run.id, "failed", str(exc))
    except client.Unreachable as exc:
        # Expected while Home Assistant restarts: the next tick asks again.
        logger.info("Home Assistant update %s: %s", run.id, exc)
        await _past_deadline(app, run)


def _elapsed(since: datetime) -> float:
    return (datetime.now(timezone.utc) - since).total_seconds()


async def _past_deadline(app, run) -> bool:
    if run.state == "backing_up" and _elapsed(run.started_at) > BACKUP_DEADLINE_S:
        await _finish(app, run.id, "failed",
                      f"No new backup within {BACKUP_DEADLINE_S // 60} minutes. Nothing was installed.")
        return True
    if run.state == "installing":
        limit = INSTALL_DEADLINE_S.get(run.kind, INSTALL_DEADLINE_S["os"])
        if _elapsed(run.install_sent_at or run.backup_at or run.started_at) > limit:
            await _finish(app, run.id, "failed",
                          f"Still not on {run.to_version} after {limit // 60} minutes. Look in Home Assistant.")
            return True
    return False


async def _advance_backup(app, run, connection, transport) -> None:
    token, _webhook = _secrets(connection)
    url, verify = connection.url, connection.verify_tls
    backup = await client.read_state(url, LAST_BACKUP, token, verify, transport)
    value = (backup or {}).get("state")
    if is_newer(value, run.backup_before):
        async with app.db_sessionmaker() as session:
            async with session.begin():
                row = await session.get(HomeAssistantUpdate, run.id)
                row.state, row.backup_at = "installing", parse_moment(value)
        return
    event = await client.read_state(url, BACKUP_EVENT, token, verify, transport)
    if event and event.get("state") != run.backup_event_before \
            and (event.get("attributes") or {}).get("event_type") == "failed":
        reason = (event.get("attributes") or {}).get("failed_reason")
        await _finish(app, run.id, "failed",
                      f"The backup failed{f' ({reason})' if reason else ''}. Nothing was installed.")
        return
    if _elapsed(run.started_at) > TRIGGER_GRACE_S:
        automation = await client.read_state(url, automation_entity_id(connection), token, verify, transport)
        triggered = ((automation or {}).get("attributes") or {}).get("last_triggered")
        if not is_newer(triggered, run.triggered_before):
            await _finish(app, run.id, "failed",
                          "Home Assistant didn't run the maintenance automation, so its webhook id may not "
                          "match Steward's, or Steward isn't on its local network. Nothing was installed.")
            return
    await _past_deadline(app, run)


async def _advance_install(app, run, connection, transport) -> None:
    token, webhook_id = _secrets(connection)
    url, verify = connection.url, connection.verify_tls
    if run.install_sent_at is None:
        await client.send_webhook(url, webhook_id, {"action": "install", "entity_id": run.entity_id},
                                  verify, transport)
        async with app.db_sessionmaker() as session:
            async with session.begin():
                row = await session.get(HomeAssistantUpdate, run.id)
                row.install_sent_at = datetime.now(timezone.utc)
        return
    current = await client.read_state(url, run.entity_id, token, verify, transport)
    attrs = (current or {}).get("attributes") or {}
    installed = attrs.get("installed_version")
    if installed and (installed == run.to_version
                      or (current.get("state") == "off" and installed != run.from_version)):
        await _finish(app, run.id, "done", None)
        return
    await _past_deadline(app, run)