arr-init/scripts/common.py

"""Shared utilities for arr-init scripts."""

import json
import sys
import time
import xml.etree.ElementTree as ET

import requests as http
import yaml


def load_config():
    """Load JSON configuration from the path given as the first CLI argument."""
    if len(sys.argv) < 2:
        print("Usage: script <config.json>", file=sys.stderr)
        sys.exit(1)
    with open(sys.argv[1]) as f:
        return json.load(f)


def read_api_key_xml(config_xml_path):
    """Extract <ApiKey> from a Servarr config.xml file."""
    tree = ET.parse(config_xml_path)
    node = tree.find("ApiKey")
    if node is None or not node.text:
        raise ValueError(f"Could not find ApiKey in {config_xml_path}")
    return node.text


def read_api_key_yaml(config_yaml_path):
    """Extract the apikey from Bazarr's config.yaml (auth section)."""
    with open(config_yaml_path) as fh:
        data = yaml.safe_load(fh)
    try:
        return data["auth"]["apikey"]
    except (KeyError, TypeError) as exc:
        raise ValueError(
            f"Could not find auth.apikey in {config_yaml_path}"
        ) from exc


def wait_for_api(
    base_url,
    api_key,
    timeout,
    name,
    *,
    header_name="X-Api-Key",
    status_path="/system/status",
):
    """Poll a status endpoint until the API responds or timeout.

    Args:
        base_url: Base URL including any API version prefix.
        api_key: API key for authentication.
        timeout: Maximum seconds to wait.
        name: Human-readable service name for log messages.
        header_name: HTTP header name for the API key.
        status_path: Path appended to base_url for the health probe.
    """
    print(f"Waiting for {name} API (timeout: {timeout}s)...")
    for i in range(1, timeout + 1):
        try:
            resp = http.get(
                f"{base_url}{status_path}",
                headers={header_name: api_key},
                timeout=5,
            )
            if resp.ok:
                print(f"{name} API is ready")
                return
        except (http.ConnectionError, http.Timeout):
            pass
        if i == timeout:
            print(
                f"{name} API not available after {timeout} seconds",
                file=sys.stderr,
            )
            sys.exit(1)
        time.sleep(1)


def health_check_loop(url, api_key, entity_name, svc_name, max_retries, interval):
    """POST to a testall endpoint with retry logic.

    Exits the process on permanent failure so the systemd unit reflects the error.
    """
    attempt = 0
    while True:
        healthy = True
        last_error = ""
        try:
            resp = http.post(
                url,
                headers={
                    "X-Api-Key": api_key,
                    "Content-Type": "application/json",
                },
                timeout=30,
            )
            result = resp.json()
            failures = [
                item for item in result if not item.get("isValid", True)
            ]
            if failures:
                healthy = False
                last_error = "\n".join(
                    f"  - ID {f['id']}: "
                    + ", ".join(
                        v["errorMessage"]
                        for v in f.get("validationFailures", [])
                    )
                    for f in failures
                )
        except (http.RequestException, ValueError, KeyError) as exc:
            healthy = False
            last_error = (
                f"could not reach {svc_name} API for {entity_name} test: {exc}"
            )

        if healthy:
            print(f"All {entity_name}s healthy")
            return

        attempt += 1
        if attempt > max_retries:
            print(
                f"Health check FAILED after {attempt} attempts: "
                f"{entity_name}(s) unreachable:",
                file=sys.stderr,
            )
            print(last_error, file=sys.stderr)
            sys.exit(1)

        print(
            f"{entity_name.capitalize()} health check failed "
            f"(attempt {attempt}/{max_retries}), "
            f"retrying in {interval}s..."
        )
        time.sleep(interval)