homeassistant-stelloauth-addon/stelloauth/rootfs/usr/local/bin/addon-supervisor

#!/usr/bin/env python3
from __future__ import annotations

import dataclasses
import json
import logging
import os
import re
import shutil
import signal
import subprocess
import sys
import time
import urllib.parse
import urllib.request
from pathlib import Path
from typing import Callable, NamedTuple


OPTIONS_PATH = Path("/data/options.json")
PROFILE_PATH = Path("/tmp/cloakserve")
CLOAK_ROOT = "http://127.0.0.1:9222/"
CLOAK_VERSION = "http://127.0.0.1:9222/json/version?fingerprint=addon-readiness"
CLOAK_CLOSE = "http://127.0.0.1:9222/fingerprint/addon-readiness/close"
STELLOAUTH_ROOT = "http://127.0.0.1:8080/"
DURATION = re.compile(r"^[1-9][0-9]*(?:ms|s|m|h)$")
CLOAK_COMMAND = [
    "/usr/local/bin/cloakserve",
    "--headless=true",
    "--idle-timeout=30",
    "--data-dir=/tmp/cloakserve",
]
STELLOAUTH_COMMAND = ["/usr/local/bin/stelloauth"]


class ConfigError(RuntimeError):
    pass


@dataclasses.dataclass(frozen=True)
class Options:
    queue_timeout: str
    rate_limit_count: int
    rate_limit_duration: str


class HttpResponse(NamedTuple):
    status: int
    body: bytes


class ReadinessError(RuntimeError):
    pass


def load_options(path: Path) -> Options:
    try:
        value = json.loads(path.read_text(encoding="utf-8"))
        if not isinstance(value, dict) or set(value) != {
            "queue_timeout",
            "rate_limit_count",
            "rate_limit_duration",
        }:
            raise ValueError

        queue_timeout = value["queue_timeout"]
        rate_limit_count = value["rate_limit_count"]
        rate_limit_duration = value["rate_limit_duration"]
        if not isinstance(queue_timeout, str) or DURATION.fullmatch(queue_timeout) is None:
            raise ValueError
        if (
            isinstance(rate_limit_count, bool)
            or not isinstance(rate_limit_count, int)
            or not 1 <= rate_limit_count <= 20
        ):
            raise ValueError
        if (
            not isinstance(rate_limit_duration, str)
            or DURATION.fullmatch(rate_limit_duration) is None
        ):
            raise ValueError
        return Options(queue_timeout, rate_limit_count, rate_limit_duration)
    except (OSError, UnicodeError, json.JSONDecodeError, KeyError, TypeError, ValueError):
        raise ConfigError("Invalid add-on configuration") from None


def build_environment(options: Options) -> dict[str, str]:
    environment = os.environ.copy()
    environment.update(
        {
            "CLOAK_CDP_URL": "http://127.0.0.1:9222",
            "CLOAK_MAX_SESSIONS": "1",
            "CLOAK_QUEUE_TIMEOUT": options.queue_timeout,
            "RATE_LIMIT_COUNT": str(options.rate_limit_count),
            "RATE_LIMIT_DURATION": options.rate_limit_duration,
            "HTTP_ADDRESS": "0.0.0.0",
            "PORT": "8080",
            "METRICS_ADDRESS": "127.0.0.1",
            "METRICS_PORT": "9090",
        }
    )
    return environment


def http_request(method: str, url: str, timeout: float) -> HttpResponse:
    opener = urllib.request.build_opener(urllib.request.ProxyHandler({}))
    data = b"" if method == "POST" else None
    request = urllib.request.Request(url, data=data, method=method)
    with opener.open(request, timeout=timeout) as response:
        return HttpResponse(response.status, response.read())


def probe_cloak(
    request: Callable[[str, str, float], HttpResponse] = http_request,
) -> None:
    if request("GET", CLOAK_ROOT, 2.0).status != 200:
        raise ReadinessError("CloakBrowser root probe failed")

    version = request("GET", CLOAK_VERSION, 2.0)
    if version.status != 200:
        raise ReadinessError("CloakBrowser CDP probe failed")
    try:
        document = json.loads(version.body)
        if not isinstance(document, dict):
            raise ValueError
        websocket_url = document["webSocketDebuggerUrl"]
        if (
            not isinstance(websocket_url, str)
            or not websocket_url.startswith("ws://127.0.0.1:")
        ):
            raise ValueError
        parsed = urllib.parse.urlsplit(websocket_url)
        if (
            parsed.scheme != "ws"
            or parsed.hostname != "127.0.0.1"
            or parsed.username is not None
            or parsed.password is not None
            or parsed.port is None
        ):
            raise ValueError
    except (json.JSONDecodeError, KeyError, TypeError, UnicodeError, ValueError):
        raise ReadinessError("CloakBrowser CDP response invalid") from None

    if request("POST", CLOAK_CLOSE, 2.0).status != 200:
        raise ReadinessError("CloakBrowser readiness profile close failed")


def probe_stelloauth(
    request: Callable[[str, str, float], HttpResponse] = http_request,
) -> None:
    if request("GET", STELLOAUTH_ROOT, 2.0).status != 200:
        raise ReadinessError("Stelloauth root probe failed")


def wait_until_ready(
    name: str,
    probe: Callable[[], None],
    timeout: float,
    stopping: Callable[[], bool],
    monotonic: Callable[[], float] = time.monotonic,
    sleep: Callable[[float], None] = time.sleep,
) -> None:
    deadline = monotonic() + timeout
    delay = 0.25
    while monotonic() < deadline and not stopping():
        try:
            probe()
            return
        except (OSError, ValueError, ReadinessError):
            sleep(delay)
            delay = min(delay * 2, 2.0)
    raise ReadinessError(f"{name} did not become ready")


class ProcessManager:
    def __init__(
        self,
        environment: dict[str, str],
        *,
        popen: Callable[..., subprocess.Popen] = subprocess.Popen,
        killpg: Callable[[int, int], None] = os.killpg,
        cloak_probe: Callable[[], None] = probe_cloak,
        stelloauth_probe: Callable[[], None] = probe_stelloauth,
        monotonic: Callable[[], float] = time.monotonic,
        sleep: Callable[[float], None] = time.sleep,
        profile_path: Path = PROFILE_PATH,
    ) -> None:
        self.environment = environment
        self._popen = popen
        self._killpg = killpg
        self._cloak_probe = cloak_probe
        self._stelloauth_probe = stelloauth_probe
        self._monotonic = monotonic
        self._sleep = sleep
        self._profile_path = profile_path
        self._children: list[tuple[str, subprocess.Popen]] = []
        self._stopping = False
        self._term_sent = False
        self._shutdown_deadline: float | None = None

    def _handle_signal(self, _signum: int, _frame: object) -> None:
        if self._stopping:
            return
        logging.info("Shutdown requested")
        self._stopping = True
        self._begin_shutdown()

    def _clean_profiles(self) -> None:
        logging.info("Cleaning CloakBrowser profiles")
        if self._profile_path.is_symlink():
            self._profile_path.unlink()
        elif self._profile_path.exists():
            shutil.rmtree(self._profile_path)
        self._profile_path.mkdir(parents=True, mode=0o700)
        self._profile_path.chmod(0o700)

    def _spawn(self, name: str, command: list[str]) -> subprocess.Popen:
        process = self._popen(
            command,
            env=self.environment,
            start_new_session=True,
        )
        self._children.append((name, process))
        if self._stopping:
            try:
                self._killpg(process.pid, signal.SIGTERM)
            except OSError:
                pass
        return process

    def _first_exited_child(self) -> tuple[str, int] | None:
        for name, process in self._children:
            returncode = process.poll()
            if returncode is not None:
                return name, returncode
        return None

    def _startup_stopping(self) -> bool:
        return self._stopping or self._first_exited_child() is not None

    def _signal_running(self, sent_signal: int) -> None:
        for _name, process in self._children:
            if process.poll() is not None:
                continue
            try:
                self._killpg(process.pid, sent_signal)
            except OSError:
                pass

    def _begin_shutdown(self) -> None:
        if self._shutdown_deadline is None:
            self._shutdown_deadline = self._monotonic() + 10.0
        if not self._term_sent:
            self._signal_running(signal.SIGTERM)
            self._term_sent = True

    def _reap_all(self) -> None:
        for _name, process in self._children:
            try:
                process.wait(timeout=None)
            except OSError:
                logging.error("Child process reap failed")

    def _shutdown(self) -> None:
        if not self._children:
            return
        logging.info("Stopping child processes")
        self._begin_shutdown()
        assert self._shutdown_deadline is not None
        while (
            any(process.poll() is None for _name, process in self._children)
            and self._monotonic() < self._shutdown_deadline
        ):
            remaining = self._shutdown_deadline - self._monotonic()
            self._sleep(min(0.25, max(0.0, remaining)))

        if any(process.poll() is None for _name, process in self._children):
            logging.info("Forcing child processes to stop")
            self._signal_running(signal.SIGKILL)
        self._reap_all()
        logging.info("Child processes stopped")

    @staticmethod
    def _failure_status(returncode: int) -> int:
        return returncode if returncode != 0 else 1

    def _startup_failure(self, service: str) -> int:
        exited = self._first_exited_child()
        if exited is not None:
            name, returncode = exited
            logging.error(f"{name} exited unexpectedly")
            status = self._failure_status(returncode)
        else:
            logging.error(f"{service} readiness failed")
            status = 1
        self._shutdown()
        return status

    def _run(self) -> int:
        try:
            self._clean_profiles()
            if self._stopping:
                self._shutdown()
                return 0

            logging.info("Starting CloakBrowser")
            self._spawn("CloakBrowser", CLOAK_COMMAND)
            wait_until_ready(
                "CloakBrowser",
                self._cloak_probe,
                60.0,
                self._startup_stopping,
                self._monotonic,
                self._sleep,
            )
            if self._stopping:
                self._shutdown()
                return 0
            if self._first_exited_child() is not None:
                return self._startup_failure("CloakBrowser")
            logging.info("CloakBrowser ready")

            logging.info("Starting Stelloauth")
            self._spawn("Stelloauth", STELLOAUTH_COMMAND)
            wait_until_ready(
                "Stelloauth",
                self._stelloauth_probe,
                30.0,
                self._startup_stopping,
                self._monotonic,
                self._sleep,
            )
            if self._stopping:
                self._shutdown()
                return 0
            if self._first_exited_child() is not None:
                return self._startup_failure("Stelloauth")
            logging.info("Stelloauth listening on 0.0.0.0:8080")
        except ReadinessError:
            if self._stopping:
                self._shutdown()
                return 0
            service = "Stelloauth" if len(self._children) > 1 else "CloakBrowser"
            return self._startup_failure(service)
        except OSError:
            logging.error("Process startup failed")
            self._shutdown()
            return 1

        while not self._stopping:
            exited = self._first_exited_child()
            if exited is not None:
                name, returncode = exited
                logging.error(f"{name} exited unexpectedly")
                status = self._failure_status(returncode)
                self._shutdown()
                return status
            self._sleep(0.25)

        self._shutdown()
        return 0

    def run(self) -> int:
        previous_handlers = {
            signal.SIGTERM: signal.signal(signal.SIGTERM, self._handle_signal),
            signal.SIGINT: signal.signal(signal.SIGINT, self._handle_signal),
        }
        try:
            return self._run()
        finally:
            for handled_signal, previous_handler in previous_handlers.items():
                signal.signal(handled_signal, previous_handler)


def main() -> int:
    logging.basicConfig(level=logging.INFO, format="%(message)s")
    try:
        options = load_options(OPTIONS_PATH)
    except ConfigError:
        logging.error("Invalid add-on configuration")
        return 2
    return ProcessManager(build_environment(options)).run()


if __name__ == "__main__":
    sys.exit(main())