archive-clients/src/archive_clients/services.py

"""Bounded, read-only startup probes for local service capabilities."""

from __future__ import annotations

import json
import re
from dataclasses import dataclass
from datetime import datetime, timezone
from http.cookiejar import CookieJar
from typing import Any
from urllib import error, parse, request

from archive_clients.config import ServiceConfig
from archive_clients.routes import discover_routes
from archive_control.v1 import common_pb2, route_pb2


_MAX_RESPONSE_BYTES = 1024 * 1024
_QB_VERSION = re.compile(r"^v?(\d+)\.(\d+)(?:\.|$)")


@dataclass(frozen=True)
class ServiceProbe:
    service: str
    state: int
    checked_at: datetime
    version: str = ""
    api_version: str = ""
    detail: str = ""
    device_id: str = ""
    libtorrent_version: str = ""
    routes: tuple[route_pb2.LocalRoute, ...] = ()


class _NoRedirect(request.HTTPRedirectHandler):
    def redirect_request(self, *_: Any, **__: Any) -> None:
        return None


def probe_qbittorrent(
    config: ServiceConfig, timeout: float = 10
) -> ServiceProbe:
    checked_at = datetime.now(timezone.utc)
    try:
        opener = request.build_opener(
            request.HTTPCookieProcessor(CookieJar()), _NoRedirect()
        )
        credentials = parse.urlencode({
            "username": config.username,
            "password": config.read_password(),
        }).encode("utf-8")
        login = request.Request(
            _url(config.endpoint, "/api/v2/auth/login"),
            data=credentials,
            method="POST",
            headers={"Content-Type": "application/x-www-form-urlencoded"},
        )
        login_response = opener.open(login, timeout=timeout)
        login_status = getattr(login_response, "status", 200)
        login_body = _read(login_response).strip()
        if login_body != "Ok." and not (
            login_status == 204 and not login_body
        ):
            raise PermissionError("qBittorrent authentication failed")
        version = _text_get(opener, config.endpoint, "/api/v2/app/version", timeout)
        api_version = _text_get(
            opener, config.endpoint, "/api/v2/app/webapiVersion", timeout
        )
        build = _json_get(
            opener, config.endpoint, "/api/v2/app/buildInfo", timeout
        )
        libtorrent = build.get("libtorrent", "")
        if not all(isinstance(value, str) and value for value in (
            version, api_version, libtorrent,
        )):
            raise ValueError("qBittorrent returned incomplete version data")
        if not _supported_qb_version(version):
            return ServiceProbe(
                "qbittorrent", common_pb2.HEALTH_STATE_UNHEALTHY,
                checked_at, version=version, api_version=api_version,
                detail="unsupported qBittorrent version",
                libtorrent_version=libtorrent,
            )
        return ServiceProbe(
            "qbittorrent", common_pb2.HEALTH_STATE_HEALTHY, checked_at,
            version=version, api_version=api_version,
            libtorrent_version=libtorrent,
        )
    except (PermissionError, error.HTTPError) as exc:
        detail = "authentication failed" if isinstance(
            exc, PermissionError
        ) or getattr(exc, "code", 0) in {401, 403} else "HTTP probe failed"
        return ServiceProbe(
            "qbittorrent", common_pb2.HEALTH_STATE_UNHEALTHY,
            checked_at, detail=detail,
        )
    except (error.URLError, OSError, UnicodeError, ValueError, json.JSONDecodeError):
        return ServiceProbe(
            "qbittorrent", common_pb2.HEALTH_STATE_DEGRADED,
            checked_at, detail="service unavailable or incompatible",
        )


def probe_syncthing(
    config: ServiceConfig,
    timeout: float = 10,
    sparse_supported: bool = False,
) -> ServiceProbe:
    checked_at = datetime.now(timezone.utc)
    try:
        opener = request.build_opener(_NoRedirect())
        headers = {"X-API-Key": config.read_api_key() or ""}
        version = _json_get(
            opener, config.endpoint, "/rest/system/version", timeout, headers
        )
        status = _json_get(
            opener, config.endpoint, "/rest/system/status", timeout, headers
        )
        configuration = _json_get(
            opener, config.endpoint, "/rest/config", timeout, headers
        )
        service_version = version.get("longVersion") or version.get("version")
        device_id = status.get("myID")
        if not isinstance(service_version, str) or not service_version:
            raise ValueError("Syncthing returned no version")
        if not isinstance(device_id, str) or not device_id:
            raise ValueError("Syncthing returned no device ID")
        routes = discover_routes(
            configuration, device_id, config.roots, sparse_supported,
            checked_at,
        )
        return ServiceProbe(
            "syncthing", common_pb2.HEALTH_STATE_HEALTHY, checked_at,
            version=service_version, device_id=device_id,
            routes=tuple(routes),
        )
    except error.HTTPError as exc:
        detail = (
            "authentication failed"
            if exc.code in {401, 403} else "HTTP probe failed"
        )
        return ServiceProbe(
            "syncthing", common_pb2.HEALTH_STATE_UNHEALTHY,
            checked_at, detail=detail,
        )
    except (error.URLError, OSError, UnicodeError, ValueError, json.JSONDecodeError):
        return ServiceProbe(
            "syncthing", common_pb2.HEALTH_STATE_DEGRADED,
            checked_at, detail="service unavailable or incompatible",
        )


def _url(endpoint: str, path: str) -> str:
    return f"{endpoint.rstrip('/')}{path}"


def _supported_qb_version(version: str) -> bool:
    match = _QB_VERSION.match(version)
    if match is None:
        return False
    major, minor = int(match.group(1)), int(match.group(2))
    # torrents/export, required to preserve exact metainfo before staging,
    # was introduced in qBittorrent 4.5.0.
    return major == 5 or (major == 4 and minor in {5, 6})


def _text_get(
    opener: Any, endpoint: str, path: str, timeout: float,
    headers: dict[str, str] | None = None,
) -> str:
    call = request.Request(_url(endpoint, path), headers=headers or {})
    return _read(opener.open(call, timeout=timeout)).strip()


def _json_get(
    opener: Any, endpoint: str, path: str, timeout: float,
    headers: dict[str, str] | None = None,
) -> dict[str, Any]:
    value = json.loads(_text_get(opener, endpoint, path, timeout, headers))
    if not isinstance(value, dict):
        raise ValueError("service response is not an object")
    return value


def _read(response: Any) -> str:
    with response:
        data = response.read(_MAX_RESPONSE_BYTES + 1)
    if len(data) > _MAX_RESPONSE_BYTES:
        raise ValueError("service response exceeds limit")
    return data.decode("utf-8")