#!/usr/bin/env bash
# ups companion for openbsd // i just didn't fancy running nut after all
# (c) 2024 th@bogus.net 
# $upsuser  must be allowed to ssh into your hosts, even when /etc/nologin is present
# add the following line to /etc/pam.d/sshd (assuming this user's group is _nut)
# account    sufficient   pam_succeed_if.so user ingroup _nut
# account    required     pam_nologin.so
# and restart sshd
# hosts must be configured with $upsuser, ssh keys, and be allowed to run /sbin/shutdown 
# as root, e.g. add something like this to /etc/sudoers
# _ups ALL=NOPASSWD:/sbin/shutdown

### start of user-configurable variables
upsuser='_nut'
notifyaddress='' # your email address
grace_time="+5" # minutes, adjust for your battery capacity
ups_kernel_variable="hw.sensors.upd0" 
default_hostnames=("host1" "host2" "host3")
daemon_mode=0
pushover=0
TOKEN="" # pushover token
USER="" # pushover user token
mail=0
syslog=0
FILENAME="/tmp/upsc.lock.79a41b38b2ccaf2692d72800c9443dea"
### end of user-configurable variables

LOCKFILE=""

## all user-configurable variables can be redefined in $HOME/.upscrc
if [ -r $HOME/.upscrc ]; then
    source $HOME/.upscrc
fi

usage() {
    if [ -z "$TOKEN" ]; then
        pushover_token="NO API TOKEN!"
    else
        pushover_token="defined"
    fi
    echo "Usage: $0 [-h] [-p] [-m] [-s] [-d] [-t <time>] [<shutdown|cancel>] [<hostname1> [<hostname2> ...]]"
    echo "  -h: this"
    echo "  -t: optional grace time (default: $grace_time)"
    echo "  -p: alert via Pushover ($pushover_token)"
    echo "  -m: alert via mail ($notifyaddress)"
    echo "  -s: disable logging to syslog"
    echo "  -d: daemon mode"
    exit 1
}

power_check() {
    output=$(sysctl $ups_kernel_variable) 

    if [ -z "$output" ]; then
        echo "Fatal: could not retrieve UPS status: $output"
        exit 1
    fi

    ac_present=$(echo "$output" | grep 'ACPresent' | grep -oE '(On|Off)')

    if [ "$ac_present" == "On" ]; then
        current_power_status="power_on"
    else
        current_power_status="power_off"
    fi
}

pushover_notify() {
    pushover_message=$1
    pushover_status=$(curl -s \
      --form-string "token=$TOKEN" \
      --form-string "user=$USER" \
      --form-string "message=$pushover_message" \
      https://api.pushover.net/1/messages.json | jq -r .status)
    if [ $pushover_status -ne 1 ];then
        echo "Warning: Pushover notification didn't succeed, server responded with: $pushover_status"
    fi
}

mail_notify() {
    mail_message=$1
    mail -s "$mail_message" $notifyaddress <<EOF
UPS status:

$(sysctl $ups_kernel_variable)
EOF
}

execute_command() {
    for hostname in "${hostnames[@]}"; do
        if [ -n "$command" ]; then
            echo "Executing '$command' on host '$hostname'"
            # Here you can add your actual shutdown or cancel command, possibly using ssh for remote hostnames
            if [ "$command" == "shutdown" ]; then
                ssh $upsuser@$hostname sudo /sbin/shutdown $grace_time
            else
                ssh $upsuser@$hostname sudo /sbin/shutdown -c 
            fi
        else
            check_ups_status $hostname
        fi
    done
    if [ "$current_power_status" != "$previous_power_status" ]; then
        if [ "$current_power_status" == "power_off" ]; then
            msg="Power loss detected. UPS is now on battery power."
        else
            msg="Power restored. UPS is back on AC power."
        fi
        echo $msg
        if [ $syslog -eq 1 ]; then
            logger $(basename $0) "$msg"
        fi
        if [ $pushover -eq 1 ];then
            pushover_notify "$msg"
        fi
        if [ $mail -eq 1 ];then
            mail_notify "$msg"
        fi
    fi
    previous_power_status="$current_power_status"
}

check_ups_status() {
    hostname=$1
    power_check
    if [ "$current_power_status" != "$previous_power_status" ]; then
        if [ "$current_power_status" == "power_off" ]; then
            ssh $upsuser@$hostname sudo /sbin/shutdown $grace_time
        else
            ssh $upsuser@$hostname sudo /sbin/shutdown -c
        fi
    fi
}

cleanup() {
    msg="Caught signal, exiting."
    echo $msg
    if [ $syslog -eq 1 ]; then
        logger $(basename $0) "$msg"
    fi
    if [ $pushover -eq 1 ];then
        pushover_notify "$msg"
    fi
    if [ $mail -eq 1 ];then
        mail_notify "$msg"
    fi    
    if [ ! -z $LOCKFILE ]; then
        rm -f $LOCKFILE
    fi
    exit
}

check_and_recreate_lockfile() {
    if [ ! -e "$LOCKFILE" ]; then
        msg="Lockfile missing, recreating."
        pushover_notify "$msg"
        mail_notify "$msg"
        echo $$ > "$LOCKFILE"
    fi
}

trap cleanup SIGINT SIGTERM SIGQUIT SIGABRT SIGHUP

while getopts ":hdpmst:" opt; do
    case ${opt} in
        h )
            usage
            ;;
        d )
            daemon_mode=1
            LOCKFILE=$FILENAME
            if [ -e "$LOCKFILE" ] && kill -0 $(cat "$LOCKFILE") 2>/dev/null; then
                echo "Another instance is already running, exiting."
                exit 0
            else
                echo $$ > "$LOCKFILE"
            fi
            ;;
        t )
            grace_time="$OPTARG"
            if ! [[ "$grace_time" =~ ^(\+[0-9]+|now)$ ]]; then
                echo "Error: Grace time must be a number or 'now'."
                usage
            fi
            ;;
        p )
            if [ -z $TOKEN ]; then
                echo "Error: no Pushover token defined in $HOME/.upscrc"
                exit 1
            fi
            if [ -z $(which curl) ]; then
                echo "Error: curl is required for Pushover notifications."
                exit 1
            fi
            if [ -z $(which jq) ]; then
                echo "Error: need jq to parse json status from Pushover server."
                exit 1
            fi
            pushover=1
            ;;
        m )
            if [ -z $notifyaddress ]; then
                echo "Error: no mail address specified"
                exit 1
            fi
            mail=1
            ;;
        s )
            syslog=0
            ;;
        \? )
            echo "Invalid Option: -$OPTARG" 1>&2
            usage
            ;;
    esac
done

shift $((OPTIND -1))

if [ $daemon_mode -eq 0 ]; then
    command=$1; shift
    case "$command" in
        shutdown|cancel)
            ;;
        *)
            echo "Invalid command: $command"
            usage
            ;;
    esac
else
    command="" # daemon mode
fi

hostnames=("$@") # remaining arguments are considered hostnames

if [ ${#hostnames[@]} -eq 0 ]; then
    hostnames=("${default_hostnames[@]}")
fi

power_check

if [ "$current_power_status" == "power_off" ]; then
    msg="Starting: UPS is on battery power."
else
    msg="Starting: UPS is on AC power."
fi

echo $msg
if [ $syslog -eq 1 ]; then
    logger $(basename $0) "$msg"
fi
if [ $pushover -eq 1 ];then
    pushover_notify "$msg"
fi
if [ $mail -eq 1 ];then
    mail_notify "$msg"
fi

previous_power_status="$current_power_status"

if [ $daemon_mode -eq 1 ]; then
    while :; do
        check_and_recreate_lockfile
        execute_command
        sleep 10
    done
else
    execute_command
fi