vacuum-wall/daemon/client.py

"""Sync daemon client for the web UI.

Communicates with vacuum-walld over a Unix socket using requests-unixsocket.
"""

import json
import logging
import re
import urllib.parse
from typing import Any

import requests
import requests_unixsocket

from daemon.iface import PathLike

logger = logging.getLogger(__name__)

_param_re = re.compile(r"<(\w+)>")


class NotFound(Exception):
    """Raised when the daemon returns HTTP 404."""

    pass


class BadRequest(Exception):
    """Exception raised when the daemon returns HTTP 400.

    Used to signal client-side input or formatting errors from the
    daemon for distinction from general failures.
    """

    pass


class Conflict(Exception):
    """Raised when the daemon returns HTTP 409."""

    pass


_DEFAULT_SOCKET = None


def _get_socket_path() -> str:
    """Return the daemon Unix socket path.

    Lazily resolves the path from the VACUUM_WALLD_SOCKET environment
    variable or falls back to data/daemon.sock under the project
    directory.  The result is cached in _DEFAULT_SOCKET.

    Returns:
        Absolute path string for the daemon socket.
    """
    global _DEFAULT_SOCKET
    if _DEFAULT_SOCKET is None:
        import os
        from pathlib import Path

        project_dir = Path(__file__).resolve().parent.parent
        socket_env = os.environ.get(
            "VACUUM_WALLD_SOCKET", str(project_dir / "data" / "daemon.sock")
        )
        _DEFAULT_SOCKET = socket_env
    return _DEFAULT_SOCKET


def set_socket_path(path: str) -> None:
    """Override the default daemon socket path.

    Useful for tests that need an alternate socket.

    Args:
        path: Absolute path to the Unix socket file.
    """
    global _DEFAULT_SOCKET
    _DEFAULT_SOCKET = path


def _format_path(path: str, params: dict[str, Any] | None) -> str:
    """Replace ``<param>`` path segments with URL-encoded values from *params*.

    Args:
        path: URL path that may contain ``<key>`` placeholders.
        params: Dict of parameter values to substitute.

    Returns:
        Path with all ``<key>`` segments replaced by their URL-encoded
        values. Unmatched placeholders are left unchanged.
    """
    if params is None:
        return path

    def _replace(m: re.Match[str]) -> str:
        key = m.group(1)
        if key in params:
            return urllib.parse.quote(str(params[key]), safe="")
        return m.group(0)

    return _param_re.sub(_replace, path)


def _resolve_path(method_or_ep: PathLike, path: str | None = None) -> tuple[str, str]:
    """Resolve method/path from an Endpoint tuple or two separate arguments."""
    if isinstance(method_or_ep, tuple):
        return (method_or_ep[0], method_or_ep[1])
    if path is None:
        raise ValueError("path is required when method is a string")
    return (method_or_ep, path)


def request(
    method: PathLike,
    path: str | None = None,
    json_body: dict[str, Any] | None = None,
    query_params: dict[str, Any] | None = None,
    socket_path: str | None = None,
    timeout: float = 30,
) -> dict[str, Any]:
    """Make a request to the daemon and return the parsed response body.

    *method* can be an :class:`Endpoint` tuple from :mod:`daemon.iface`,
    in which case *path* should be omitted.

    For GET requests, query_params are sent as URL query parameters instead
    of a JSON body. For other methods, json_body is sent as JSON.

    Raises RuntimeError on non-2xx responses or connection errors.
    Raises NotFound on HTTP 404. Raises BadRequest on HTTP 400.
    """
    resolved_method, resolved_path = _resolve_path(method, path)

    # Substitute <param> segments from body/query params so the daemon
    # receives a concrete path instead of a template.
    # Merge body and query params for <param> substitution. query_params
    # takes precedence on key conflicts, so callers should avoid passing
    # the same key in both dicts.
    combined = {**(json_body or {}), **(query_params or {})}
    formatted_path = _format_path(resolved_path, combined)

    sp = socket_path or _get_socket_path()
    url = f"http+unix://{urllib.parse.quote(sp, safe='')}{formatted_path}"
    sess = requests_unixsocket.Session()
    try:
        kwargs: dict[str, Any] = {
            "timeout": timeout,
        }
        if resolved_method == "GET":
            if query_params:
                kwargs["params"] = query_params
        else:
            if json_body is not None:
                kwargs["json"] = json_body
        resp = sess.request(
            resolved_method,
            url,
            **kwargs,
        )
        resp.raise_for_status()
        try:
            data = resp.json()
        except json.JSONDecodeError as exc:
            raise RuntimeError(
                f"Daemon returned non-JSON response ({resp.status_code}): {exc}"
            ) from exc
    except requests.ConnectionError as exc:
        raise RuntimeError(f"Cannot connect to daemon at {sp}: {exc}") from exc
    except requests.Timeout as exc:
        raise RuntimeError(
            f"Daemon request timed out: {resolved_method} {resolved_path}"
        ) from exc
    except requests.HTTPError as exc:
        try:
            data = resp.json()
        except Exception:
            data = {"ok": False, "error": resp.text}
        if not data.get("ok"):
            if resp.status_code == 404:
                raise NotFound(data.get("error", str(exc))) from exc
            if resp.status_code == 400:
                raise BadRequest(data.get("error", str(exc))) from exc
            if resp.status_code == 409:
                raise Conflict(data.get("error", str(exc))) from exc
            raise RuntimeError(data.get("error", str(exc))) from exc
    if not data.get("ok"):
        raise RuntimeError(data.get("error", "Unknown error"))
    return data.get("data")


def get(path: PathLike, params: dict[str, Any] | None = None, **kwargs: Any) -> Any:
    """Send a GET request to the daemon.

    *path* can be a :class:`Endpoint` tuple from :mod:`daemon.iface`
    (e.g., ``GET_FIREWALL_ZONES``), or a plain string path.
    """
    return request(path, query_params=params, **kwargs)


def post(path: PathLike, body: dict[str, Any] | None = None, **kwargs: Any) -> Any:
    """Send a POST request to the daemon.

    The body is transmitted as a JSON payload. Extra keyword arguments
    are forwarded to request(). *path* can be an :class:`Endpoint` tuple.
    """
    return request(path, json_body=body, **kwargs)


def patch(path: PathLike, body: dict[str, Any] | None = None, **kwargs: Any) -> Any:
    """Send a PATCH request to the daemon.

    The body is transmitted as a JSON payload. Extra keyword arguments
    are forwarded to request(). *path* can be an :class:`Endpoint` tuple.
    """
    return request(path, json_body=body, **kwargs)


def delete(path: PathLike, body: dict[str, Any] | None = None, **kwargs: Any) -> Any:
    """Send a DELETE request to the daemon.

    The body is transmitted as a JSON payload. Extra keyword arguments
    are forwarded to request(). *path* can be an :class:`Endpoint` tuple.
    """
    return request(path, json_body=body, **kwargs)


def batch(ops: list[dict[str, Any]], **kwargs: Any) -> dict[str, Any]:
    """Batch request to daemon.

    Each op is a dict with 'id', 'method', 'path', and optionally 'body'.
    Returns a dict mapping each id to its result.
    """
    return request("POST", "/batch", json_body={"ops": ops}, **kwargs)