#!/bin/bash
# Nagios/Icinga plugin to check power consumption and health of a server's
# PSUs via PMBus, read in-band through the BMC (IPMI Master Write-Read).
#
# Copyright (C) 2026 Thomas Wagner <thomas.wagner@jku.at>
# SPDX-License-Identifier: GPL-2.0-or-later

VERSION="0.2"

# PMBus commands used (all reads):
#   0x79 STATUS_WORD         -> fault flags (CRITICAL / WARNING)
#   0x88 READ_VIN            -> input voltage (perfdata)
#   0x8D READ_TEMPERATURE_1  -> PSU temperature (perfdata)
#   0x90 READ_FAN_SPEED_1    -> PSU fan speed (perfdata)
#   0x97 READ_PIN            -> input power (summed, perfdata)

BUS=0x07
ADDRS="0xb0 0xb2"
WARN=0
CRIT=0
TIMEOUT=10
IPMITOOL=/usr/bin/ipmitool

usage() {
    cat <<EOF
Usage: $0 [-b bus] [-a "addr ..."] [-w watts] [-c watts] [-t seconds] [-i ipmitool] [-V] [-h]
  -b  I2C bus (default $BUS)
  -a  space-separated 8-bit PSU addresses (default "$ADDRS")
  -w  warning threshold for total power in W, 0 = none (default 0)
  -c  critical threshold for total power in W, 0 = none (default 0)
  -t  timeout per ipmitool call in seconds (default $TIMEOUT)
  -i  path of ipmitool (default $IPMITOOL)
  -V  show the version
  -h  show this help
EOF
}

unknown() { echo "UNKNOWN - $*"; exit 3; }

while getopts "b:a:w:c:t:i:Vh" opt; do
    case $opt in
        b) BUS=$OPTARG ;;
        a) ADDRS=$OPTARG ;;
        w) WARN=$OPTARG ;;
        c) CRIT=$OPTARG ;;
        t) TIMEOUT=$OPTARG ;;
        i) IPMITOOL=$OPTARG ;;
        V) echo "check_psu_power $VERSION"; exit 0 ;;
        h) usage; exit 0 ;;
        *) usage; exit 3 ;;
    esac
done

# Bus and addresses are normalized to lowercase two-digit hex ("0x7a"), the form
# the sudoers rule has to use.
hexbyte() { [[ $1 =~ ^0[xX][0-9a-fA-F]{1,2}$ ]] && printf '0x%02x' "$(( $1 ))"; }
bus=$(hexbyte "$BUS") || unknown "invalid bus '$BUS' (expected e.g. 0x07)"
BUS=$bus
addrs=()
for a in $ADDRS; do
    h=$(hexbyte "$a") || unknown "invalid PSU address '$a' (expected e.g. 0xb0)"
    addrs+=("$h")
done
[ ${#addrs[@]} -gt 0 ] || unknown "no PSU addresses given"
[[ $WARN =~ ^[0-9]+(\.[0-9]+)?$ ]] || unknown "invalid warning threshold '$WARN'"
[[ $CRIT =~ ^[0-9]+(\.[0-9]+)?$ ]] || unknown "invalid critical threshold '$CRIT'"
[[ $TIMEOUT =~ ^[0-9]+$ ]] && [ "$TIMEOUT" -gt 0 ] || unknown "invalid timeout '$TIMEOUT'"

# Format a number like the Windows variant: one decimal at most, no "-0"
fmt() { awk -v v="$1" 'BEGIN { s = sprintf("%.1f", v); sub(/\.0$/, "", s); if (s == "-0") s = "0"; print s }'; }

join() { local sep=$1 out=$2; shift 2; for x in "$@"; do out+="$sep$x"; done; printf '%s' "$out"; }

# pmbus_word ADDR CMD -> prints the 16-bit value (decimal); on failure prints the reason
pmbus_word() {
    local out rc lo hi
    out=$(timeout "$TIMEOUT" sudo -n "$IPMITOOL" raw 0x06 0x52 "$BUS" "$1" 0x02 "$2" 2>&1); rc=$?
    if [ $rc -ne 0 ]; then
        if [ $rc -eq 124 ]; then echo "timeout after ${TIMEOUT}s"
        elif [[ $out =~ rsp=(0x[0-9a-fA-F]+) ]]; then echo "completion code ${BASH_REMATCH[1],,}"
        else grep -m1 . <<< "$out" || echo "ipmitool failed"; fi
        return 1
    fi
    read -r lo hi <<< "$out"
    [[ $lo =~ ^[0-9a-fA-F]{2}$ && $hi =~ ^[0-9a-fA-F]{2}$ ]] || { echo "unexpected response '$out'"; return 1; }
    echo $((16#$hi$lo))
}

# PMBus linear11: 5-bit signed exponent, 11-bit signed mantissa
linear11() {
    awk -v raw="$1" 'BEGIN {
        e = int(raw / 2048); if (e > 15) e -= 32
        m = raw % 2048;      if (m > 1023) m -= 2048
        print m * 2 ^ e }'
}

state=0
raise() { [ "$1" -gt "$state" ] && state=$1; }

total=0
details=()
problems=()
perf=()
i=0

for addr in "${addrs[@]}"; do
    i=$((i + 1))
    psu="psu$i"

    if ! raw=$(pmbus_word "$addr" 0x97); then
        problems+=("$psu ($addr) not answering: $raw")
        raise 1
        continue
    fi
    watts=$(fmt "$(linear11 "$raw")")
    total=$(awk -v a="$total" -v b="$watts" 'BEGIN { print a + b }')
    details+=("$psu $watts W")
    perf+=("'$psu'=${watts}W")

    # STATUS_WORD: bit:severity:name. Unlisted bits (BUSY, CML, MFR, OTHER, UNKNOWN,
    # NONE_OF_THE_ABOVE) are ignored as noise; PSUs set CML after any unsupported command.
    if sw=$(pmbus_word "$addr" 0x79); then
        flags=()
        for spec in 13:2:input-fault 3:2:input-undervoltage 15:2:output-voltage-fault \
                    14:2:output-current/power-fault 4:2:overcurrent 5:2:overvoltage \
                    6:2:off 11:2:power-not-good 2:1:overtemperature 10:1:fan-fault; do
            IFS=: read -r bit sev name <<< "$spec"
            if (( (sw >> bit) & 1 )); then flags+=("$name"); raise "$sev"; fi
        done
        [ ${#flags[@]} -gt 0 ] && problems+=("$psu: ${flags[*]}")
    fi

    if raw=$(pmbus_word "$addr" 0x88); then perf+=("'${psu}_vin'=$(fmt "$(linear11 "$raw")")V"); fi
    if raw=$(pmbus_word "$addr" 0x8d); then perf+=("'${psu}_temp'=$(fmt "$(linear11 "$raw")")"); fi
    if raw=$(pmbus_word "$addr" 0x90); then perf+=("'${psu}_fan'=$(fmt "$(linear11 "$raw")")"); fi
done

if [ ${#details[@]} -eq 0 ]; then
    unknown "no PSU answered on PMBus ($(join "; " "${problems[@]}"))"
fi
total=$(fmt "$total")

ge() { awk -v a="$1" -v b="$2" 'BEGIN { exit !(a >= b) }'; }
set_() { awk -v v="$1" 'BEGIN { exit !(v > 0) }'; }   # threshold given (0 = none)
if set_ "$CRIT" && ge "$total" "$CRIT"; then raise 2
elif set_ "$WARN" && ge "$total" "$WARN"; then raise 1
fi

wt=""; set_ "$WARN" && wt=$(fmt "$WARN")
ct=""; set_ "$CRIT" && ct=$(fmt "$CRIT")
labels=(OK WARNING CRITICAL)
prefix=""
[ ${#problems[@]} -gt 0 ] && prefix="$(join "; " "${problems[@]}") - "

echo "${labels[$state]} - ${prefix}power consumption $total W (PSU PMBus: $(join ", " "${details[@]}")) | 'power'=${total}W;$wt;$ct;0; ${perf[*]}"
exit $state
