bmspy/bmspy/client.py

#
# Library with socket client for use by consumers
#
import sys
import struct
import json
import socket
from typing import Any

from bmspy.utilities import debugger
from bmspy.classes import UPS


is_registered = False


def handle_registration(socket_path: str, client_name: str, debug: int = 0) -> dict[str, Any]:
    """Register or deregister a named client with the bmspy daemon over the socket."""
    global is_registered
    data: dict[str, Any] = dict()

    if is_registered:
        message: dict[str, Any] = {"command": "DEREGISTER", "client": client_name}
    else:
        # fork server if it's not already running
        message = {"command": "REGISTER", "client": client_name}

    try:
        data = socket_comms(socket_path, message, debug)
        if data["status"] == "REGISTERED":
            is_registered = True
        elif data["status"] == "DEREGISTERED":
            is_registered = False
        else:
            raise OSError(
                "{} registration: invalid response: {}".format(client_name, data)
            )

    except Exception as e:
        if is_registered:
            debugger("{}: failed to register with daemon: {}".format(client_name, e))
        else:
            debugger("{}: failed to deregister with daemon: {}".format(client_name, e))

    return data


def _recv_exact(sock: socket.socket, length: int) -> bytes:
    """Read a complete frame, even when the stream splits it across receives."""
    data = bytearray()
    while len(data) < length:
        chunk = sock.recv(length - len(data))
        if not chunk:
            raise OSError("Daemon closed the socket before the response was complete")
        data.extend(chunk)
    return bytes(data)


def socket_comms(socket_path: str, request_data: dict[str, Any], debug: int = 0) -> dict[str, Any]:
    """Send a JSON request to the daemon socket and return the decoded JSON response."""
    response_data: dict[str, Any] = dict()

    # Create a UDS socket
    sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)

    # Connect the socket to the port where the server is listening
    if debug > 2:
        debugger("socket client: connecting to {}".format(socket_path))
    try:
        sock.connect(socket_path)
    except socket.error as msg:
        if msg.errno == 2:
            debugger("Failed to connect to bmspy daemon")
        else:
            debugger("socket client: {}".format(msg))
        sock.close()
        raise

    # Send request
    if debug > 2:
        debugger("socket client: sending {!r}".format(request_data))
    request = bytes()
    try:
        request = json.dumps(request_data).encode()
        # add length to the start of the json string, so we know how much to read on the other end
        length = struct.pack("!I", len(request))
        if debug > 3:
            debugger(
                "socket client: outgoing request length: {}, encoded as {}".format(
                    len(request), length
                )
            )
        request = length + request
        if debug > 4:
            debugger("socket client: outgoing request: {}".format(request))
    except Exception:
        sock.close()
        debugger("socket client ERROR: unable to encode request")
        sys.exit(1)

    # get length of expected json string
    try:
        sock.sendall(request)
        response = _recv_exact(sock, struct.calcsize("!I"))
        length = struct.unpack("!I", response)[0]
        if debug > 4:
            debugger(
                "socket client: incoming length: {}, encoded as {}".format(
                    length, response
                )
            )
        # read length bytes
        response = _recv_exact(sock, length)
        if debug > 3:
            debugger("socket client: incoming response: {}".format(response))
        response_data = json.loads(response)
    except Exception:
        debugger("socket client ERROR: unable to decode response")
        sys.exit(1)
    finally:
        sock.close()
    if debug > 2:
        debugger("socket client: received {!r}".format(response_data))

    return response_data


def read_data(socket_path: str, client_name: str, ups: str | None = None, debug: int = 0) -> dict[str, UPS]:
    """Return {ups_name: UPS} for all devices, or just the named one."""
    request: dict[str, Any] = {"command": "GET", "client": client_name}
    if ups is not None:
        request["ups"] = ups

    data = socket_comms(socket_path, request, debug)

    if data is None:
        raise RuntimeError("No data received from daemon")

    return {name: UPS.from_dict(fields) for name, fields in data.items()}


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