managed-container/ocf/loxthelion/managed-container.d/proxy_client.py

#!/usr/bin/env python3

import json
import socket
import sys


PROTOCOL_VERSION = 1
VALID_OCF_RCS = {0, 1, 7, 8}


class TransportFailure(Exception):
    pass


class ProtocolFailure(Exception):
    pass


def send_request(
    socket_path,
    resource,
    action,
    parameters=None,
):
    request = {
        "protocol_version": PROTOCOL_VERSION,
        "action": action,
        "resource": resource,
    }

    if parameters is not None:
        request["parameters"] = parameters

    try:
        sock = socket.socket(
            socket.AF_UNIX,
            socket.SOCK_STREAM,
        )
        sock.settimeout(2.0)

        try:
            sock.connect(socket_path)

            sock.sendall(
                json.dumps(
                    request,
                    separators=(",", ":"),
                ).encode("utf-8")
                + b"\n"
            )

            raw = b""

            while b"\n" not in raw:
                chunk = sock.recv(65536)

                if not chunk:
                    break

                raw += chunk

        finally:
            sock.close()

    except (
        socket.timeout,
        ConnectionError,
        OSError,
    ) as exc:
        raise TransportFailure from exc

    return raw


def decode_response(raw):
    if not raw:
        raise ProtocolFailure

    try:
        response = json.loads(
            raw.decode("utf-8").strip()
        )

    except (
        UnicodeDecodeError,
        json.JSONDecodeError,
    ) as exc:
        raise ProtocolFailure from exc

    if not isinstance(response, dict):
        raise ProtocolFailure

    return response


def validate_response(response, resource, action):
    if response.get("protocol_version") != PROTOCOL_VERSION:
        raise ProtocolFailure

    if response.get("ok") is not True:
        raise ProtocolFailure

    if response.get("resource") != resource:
        raise ProtocolFailure

    if response.get("action") != action:
        raise ProtocolFailure

    if "result" not in response:
        raise ProtocolFailure

    return response["result"]


def format_monitor_result(result):
    if not isinstance(result, dict):
        raise ProtocolFailure

    if result.get("delegated") is True:
        backend = result.get("backend")
        runtime = result.get("container_runtime")

        if backend != "container":
            raise ProtocolFailure

        if runtime not in ("podman", "docker"):
            raise ProtocolFailure

        return "|".join([
            "delegate",
            backend,
            runtime,
        ])

    ha = result.get("ha")

    if not isinstance(ha, dict):
        raise ProtocolFailure

    try:
        ocf_rc = int(ha.get("ocf_rc"))

    except (
        TypeError,
        ValueError,
    ) as exc:
        raise ProtocolFailure from exc

    if ocf_rc not in VALID_OCF_RCS:
        raise ProtocolFailure

    receiver = result.get("wal_receiver") or {}

    if not isinstance(receiver, dict):
        raise ProtocolFailure

    return "|".join([
        str(ocf_rc),
        str(ha.get("health", "unknown")),
        str(ha.get("ha_role", "unknown")),
        str(ha.get("state", "unknown")),
        str(
            ha.get(
                "promotion_eligible",
                "unknown",
            )
        ).lower(),
        str(result.get("role", "unknown")),
        str(
            result.get(
                "streaming",
                "unknown",
            )
        ).lower(),
        str(receiver.get("slot_name") or ""),
        str(receiver.get("sender_host") or ""),
    ])


def format_result(action, result):
    if action == "monitor":
        return format_monitor_result(result)

    return json.dumps(
        result,
        separators=(",", ":"),
        sort_keys=True,
    )


def main():
    if len(sys.argv) not in (6, 7):
        raise SystemExit(11)

    socket_path = sys.argv[1]
    resource = sys.argv[2]
    action = sys.argv[3]
    request_parameters = None

    if len(sys.argv) == 7:
        try:
            request_parameters = json.loads(sys.argv[6])

        except (
            TypeError,
            json.JSONDecodeError,
        ):
            raise SystemExit(11)

        if not isinstance(request_parameters, dict):
            raise SystemExit(11)

    try:
        transport_failure = int(sys.argv[4])
        protocol_failure = int(sys.argv[5])

    except ValueError:
        raise SystemExit(11)

    if not resource or not action:
        raise SystemExit(protocol_failure)

    try:
        raw = send_request(
            socket_path,
            resource,
            action,
            request_parameters,
        )

    except TransportFailure:
        raise SystemExit(transport_failure)

    try:
        response = decode_response(raw)
        result = validate_response(
            response,
            resource,
            action,
        )
        output = format_result(
            action,
            result,
        )

    except ProtocolFailure:
        raise SystemExit(protocol_failure)

    print(output)


if __name__ == "__main__":
    main()