bmspy/bmspy/ups.py

import argparse
import atexit
import os
import sys
import time
import smtplib
import ssl
import socket
from collections import deque
from collections.abc import Collection

from bmspy import client
from bmspy.classes import UPS, BMSScalarField

DAEMON_UPDATE_PERIOD = 30
STARTUP_RETRY_PERIOD = 5
USAGE_ERROR_STATUS = 2

scheduled_shutdown = False
critical_sent = False
warning_sent = False
alert_sent = False


def _get_field_value(device: UPS, field_name: str) -> float | None:
    """Return the raw_value of a named BMSScalarField from a UPS, or None if not found."""
    for name, field in device.items():
        if name == field_name and isinstance(field, BMSScalarField):
            return field.raw_value
    return None


def _resolve_ups_device(data: Collection[str], requested: str | None) -> str | None:
    """Return the device name to monitor, or None and print an error if the selection is invalid.

    *data* holds the available device names. A dict keyed by name also works.

    Rules:
    - No devices → error.
    - requested is set and found → use it.
    - requested is set but not found → error listing available devices.
    - requested is None and exactly one device → use it.
    - requested is None and multiple devices → error; user must specify one.
    """
    if not data:
        print("Error: no UPS devices found")
        return None
    if requested is not None:
        if requested not in data:
            print(
                "Error: UPS '{}' not found. Available: {}".format(
                    requested, ", ".join(sorted(data))
                )
            )
            return None
        return requested
    if len(data) > 1:
        print(
            "Error: multiple UPS devices found, specify one with --device: {}".format(
                ", ".join(sorted(data))
            )
        )
        return None
    return next(iter(data))


def _wait_for_ups_device(socket_path: str, requested: str | None) -> str | None:
    """Select the device to monitor, then wait until the server has data for it.

    The selection uses the configured device names, not the devices that have
    data now. A device that misses its first read at startup then cannot change
    the selection. A wrong or missing choice is final: it returns None.
    """
    name = _resolve_ups_device(client.list_devices(socket_path, "ups"), requested)
    if name is None:
        return None
    waiting = False
    while name not in client.read_data(socket_path, "ups", ups=name):
        if not waiting:
            print("UPS '{}' has no data yet, waiting...".format(name))
            waiting = True
        time.sleep(STARTUP_RETRY_PERIOD)
    return name


def handle_shutdown(action: str = "cancel", delay: int = 0, debug: int = 0) -> None:
    """Schedule or cancel a system shutdown; only schedules once per incident."""
    global scheduled_shutdown

    if action == "shutdown":
        if scheduled_shutdown is False:
            scheduled_shutdown = time.time() + delay * 60
            os.system("/sbin/shutdown {}".format(delay))

    elif action == "cancel":
        os.system("/sbin/shutdown -c")
        # Allow the next incident to schedule a shutdown again.
        scheduled_shutdown = False

    return


def handle_email(
    text: str,
    level: str | None,
    recipient: str = "root",
    mailserver: str = "localhost",
    port: int = 25,
    mailuser: str | None = None,
    mailpass: str | None = None,
    debug: int = 0,
) -> None:
    """Send an alert email using SMTP, with optional STARTTLS and authentication."""
    isSSL = False
    hostname = socket.gethostname()

    if "@" not in recipient:
        recipient = "{}@{}".format(recipient, hostname)
    sender = "bmspy on {} <nobody@{}>".format(hostname, hostname)

    if port == 465 or port == 587:
        isSSL = True

    if level is not None:
        msg = "From: {}\r\nTo: {}\r\nSubject: {} from BMSPY UPS on {}\r\n\r\n{}\r\n".format(
            sender, recipient, level, hostname, text
        )

    if isSSL:
        context = ssl.create_default_context()

    with smtplib.SMTP(mailserver, port) as server:
        if isSSL:
            server.starttls(context=context)
        server.ehlo()
        if mailuser is not None and mailpass is not None:
            server.login(mailuser, mailpass)
        server.sendmail(sender, recipient, msg)

    return


def main() -> None:
    global alert_sent, warning_sent, critical_sent

    parser = argparse.ArgumentParser(
        description="Query JBD BMS and alert or shutdown when certain thresholds are reached",
        add_help=True,
    )
    parser.add_argument(
        "--alert",
        "-a",
        dest="alert",
        action="store_true",
        default=True,
        help="Email an alert when UPS detects A/C loss (default: true)",
    )
    parser.add_argument(
        "--warning",
        "-w",
        dest="warning_threshold",
        action="store",
        type=int,
        default=75,
        help="Email an alert when remaining capacity percentage drops below this figure (default: 75)",
    )
    parser.add_argument(
        "--critical",
        "-c",
        dest="critical_threshold",
        action="store",
        type=int,
        default=30,
        help="Shut system down when remaining capacity percentage drops below this figure (default: 30)",
    )
    parser.add_argument(
        "--delay",
        "-d",
        dest="shutdown_delay",
        action="store",
        type=int,
        default=5,
        help="Delay system shutdown (default: 5 minutes)",
    )
    parser.add_argument(
        "--mailserver",
        "-m",
        dest="mailserver",
        action="store",
        default="localhost",
        help="Mail server (default: localhost)",
    )
    parser.add_argument(
        "--port",
        "-p",
        dest="port",
        action="store",
        type=int,
        default=25,
        help="Mail server port (default: 25)",
    )
    parser.add_argument(
        "--user", dest="mailuser", action="store", default=None, help="Mail server user"
    )
    parser.add_argument(
        "--pass",
        dest="mailpass",
        action="store",
        default=None,
        help="Mail server password",
    )
    parser.add_argument(
        "--to",
        "-t",
        dest="recipient",
        action="store",
        default="root",
        help="Email recipient (default: root)",
    )
    parser.add_argument(
        "--socket",
        "-s",
        dest="socket",
        action="store",
        default="/run/bmspy/bms",
        help="Socket to communicate with daemon",
    )
    parser.add_argument(
        "--device",
        dest="device",
        action="store",
        default=None,
        metavar="NAME",
        help="Name of the UPS device to monitor (required if multiple devices are available)",
    )
    parser.add_argument(
        "--verbose",
        "-v",
        action="count",
        default=0,
        help="Print more verbose information (can be specified multiple times)",
    )
    args = parser.parse_args()

    debug = args.verbose

    print("Running BMS UPS daemon on socket {}".format(args.socket))

    client.handle_registration(args.socket, "ups", debug)
    atexit.register(client.handle_registration, args.socket, "ups", debug)

    device_name = _wait_for_ups_device(args.socket, args.device)
    if device_name is None:
        # A wrong or missing --device is a usage error, as in argparse. A retry
        # cannot fix it, so the service file does not restart on this status.
        sys.exit(USAGE_ERROR_STATUS)

    history = deque()
    while True:
        data = client.read_data(args.socket, "ups", ups=device_name)
        device_obj = data.get(device_name)
        if device_obj is None:
            time.sleep(DAEMON_UPDATE_PERIOD)
            continue
        history.append(device_obj)

        # Remove the oldest data from the history
        while len(history) > 10:
            history.popleft()

        if len(history) > 3:
            comparison_1 = history[1]
            comparison_2 = history[2]
            comparison_3 = history[3]
        else:
            if debug > 1:
                print("Not enough readings, sleeping...")
            time.sleep(DAEMON_UPDATE_PERIOD)
            continue

        current_amps = float(_get_field_value(device_obj, "bms_current_amps") or 0.0)
        charge_ratio = (
            float(_get_field_value(device_obj, "bms_capacity_charge_ratio") or 0.0) * 100
        )
        comparison_1_current_amps = float(
            _get_field_value(comparison_1, "bms_current_amps") or 0.0
        )
        comparison_1_charge_ratio = (
            float(_get_field_value(comparison_1, "bms_capacity_charge_ratio") or 0.0)
            * 100
        )
        comparison_2_current_amps = float(
            _get_field_value(comparison_2, "bms_current_amps") or 0.0
        )
        comparison_2_charge_ratio = (
            float(_get_field_value(comparison_2, "bms_capacity_charge_ratio") or 0.0)
            * 100
        )
        comparison_3_current_amps = float(
            _get_field_value(comparison_3, "bms_current_amps") or 0.0
        )
        comparison_3_charge_ratio = (
            float(_get_field_value(comparison_3, "bms_capacity_charge_ratio") or 0.0)
            * 100
        )

        if debug > 1:
            print(
                "current:            {:>3.2f}A\ncapacity remaining: {:>4.0f}%".format(
                    current_amps, charge_ratio
                )
            )

        if (
            charge_ratio <= args.critical_threshold
            and comparison_1_charge_ratio <= args.critical_threshold
        ):
            if debug > 0:
                print("Below critical threshold, shutting down")
            handle_shutdown(action="shutdown", delay=args.shutdown_delay, debug=debug)
            if critical_sent is False:
                handle_email(
                    text="remaining capacity below {}%, shutting down".format(
                        args.critical_threshold
                    ),
                    level="Critical alert",
                    recipient=args.recipient,
                    mailserver=args.mailserver,
                    port=args.port,
                    mailuser=args.mailuser,
                    mailpass=args.mailpass,
                    debug=debug,
                )
                critical_sent = True

        elif (
            charge_ratio <= args.warning_threshold
            and comparison_1_charge_ratio <= args.warning_threshold
        ):
            if debug > 0:
                print("Below warning threshold")
            if warning_sent is False:
                handle_email(
                    text="remaining capacity below {}%".format(args.warning_threshold),
                    level="Warning",
                    recipient=args.recipient,
                    mailserver=args.mailserver,
                    port=args.port,
                    mailuser=args.mailuser,
                    mailpass=args.mailpass,
                    debug=debug,
                )
                warning_sent = True

        # Current needs to be negative for two consecutive reads
        elif (
            args.alert
            and current_amps < 0
            and comparison_1_current_amps < 0
            and comparison_2_current_amps >= 0
        ):
            if debug > 0:
                print("Alert: discharging!")
            if alert_sent is False:
                handle_email(
                    text="power lost",
                    level="Power loss alert",
                    recipient=args.recipient,
                    mailserver=args.mailserver,
                    port=args.port,
                    mailuser=args.mailuser,
                    mailpass=args.mailpass,
                    debug=debug,
                )
                alert_sent = True

        # Current needs to be zero or positive for two consecutive reads
        elif (
            args.alert
            and current_amps >= 0
            and comparison_1_current_amps >= 0
            and comparison_2_current_amps < 0
        ):
            if debug > 0:
                print("Alert: power regained!")
            handle_shutdown(action="cancel", debug=debug)
            handle_email(
                text="power regained",
                level="Recovery alert",
                recipient=args.recipient,
                mailserver=args.mailserver,
                port=args.port,
                mailuser=args.mailuser,
                mailpass=args.mailpass,
                debug=debug,
            )
            critical_sent = False
            warning_sent = False
            alert_sent = False

        time.sleep(DAEMON_UPDATE_PERIOD)


if __name__ == "__main__":
    main()