#!/bin/bash
# -*- sh -*-

set -e

: << =cut

=head1 NAME

wireguard_ - Wildcard-plugin to monitor wireguard peer count and traffic

=head1 CONFIGURATION

The following environment variables are used by this plugin

 active_threshold_m       - threshold to count the connection as inactive (default 3 minutes)

The plugin needs to run as root to be able to call the wg show
command.  This is configured like this:

  [wireguard_*]
      user root

This is a wildcard plugin which by default monitors all wireguard
interfaces. To monitor a single wireguard interface, link
wireguard_<interface> to this file. For example,

  ln -s /usr/share/munin/plugins/wireguard_ \
        /etc/munin/plugins/wireguard_wg0

will monitor wg0.


=head1 AUTHOR

Original author unknown

Copyright (C) 2024 pimlie

=head1 LICENSE

MIT

=head1 MAGIC MARKERS

 #%# family=auto
 #%# capabilities=autoconf suggest

=cut

. "$MUNIN_LIBDIR/plugins/plugin.sh"

INTERFACE=${0##*wireguard_}

function wg_exists {
    command -v wg >/dev/null
    return $?
}

_MAIN_ARG_FIRST=$1
function should_emit_config {
    if [ "$_MAIN_ARG_FIRST" = "config" ]; then
        return 0
    fi

    return 1
}

function should_emit_values {
    if ! should_emit_config || [ "$MUNIN_CAP_DIRTYCONFIG" = "1" ]; then
        return 0
    fi

    return 1
}

declare -a _CACHE_INTERFACES
function wg_interfaces {
    local -n _wg_ifaces=$1
    local show_all=$2

    if [ "${#_CACHE_INTERFACES[@]}" -eq 0 ]; then
        IFS=' ' read -ra _CACHE_INTERFACES <<< "$(wg show interfaces)"
    fi

    local iface
    for iface in "${_CACHE_INTERFACES[@]}"; do
        # Filter interfaces if needed
        if [ -z "$show_all" ] \
            && [ -n "$INTERFACE" ] \
            && [ "$INTERFACE" != "$iface" ]; then
            continue
        fi
        _wg_ifaces+=("$iface")
    done
}

declare -A _CACHE_PEERS
function wg_peers {
    local -n _wg_peers=$1
    local iface=$2

    if [ -z "${_CACHE_PEERS[$iface]}" ]; then
        # From wg 8 manpage:
        # If dump is specified, then several lines are printed; the first contains
        # in order separated by tab: private-key, public-key, listen-port, fwmark.
        # Subsequent lines are printed for each peer and contain in order separated
        # by tab: public-key, preshared-key, endpoint, allowed-ips, latest-handshake,
        # transfer-rx, transfer-tx, persistent-keepalive
        # Pipe to tail to skip first line
        _CACHE_PEERS["$iface"]="$(wg show "$iface" dump | tail -n +2)"
    fi

    readarray -t peers <<< "${_CACHE_PEERS[$iface]}"
    _wg_peers=("${peers[@]}")
}

function safe_peer_id {
    unsafe_peer_id=$1

    echo "${unsafe_peer_id//[^a-zA-Z0-9-_]/_}"
}

function peer_count {
    declare -a ifaces
    wg_interfaces ifaces

    echo "multigraph wireguard_peercount"

    if should_emit_config; then
        # Config for peer count per interface graph
        cat << EOF
graph_title interface peer count
graph_vlabel Number of peers
graph_category wireguard
graph_info This graph shows the number of peers per wireguard interface
EOF
    fi

    for iface in "${ifaces[@]}"; do
        if should_emit_config; then
            # List config for all interfaces
            cat <<EOF
pc_on_$iface.label $iface
pc_on_$iface.info Configured peers on interface $iface
pc_on_$iface.min 0
apc_on_$iface.label Active on $iface
apc_on_$iface.info Active peers on interface $iface
apc_on_$iface.min 0
EOF
        fi

        if should_emit_values; then
            active_threshold="${active_threshold:-$(date --date="${active_threshold_m:-3} min ago" +%s)}"

            declare -a peers
            wg_peers peers "$iface"

            peer_count="${#peers[@]}"
            active_peer_count=$(awk -F"\t" -v threshold="$active_threshold" '$5 > threshold' <<< "$(IFS=$'\n'; echo "${peers[*]}")" | wc -l)

            echo "pc_on_$iface.value $peer_count"
            echo "apc_on_$iface.value $active_peer_count"
        fi
    done
}

function peer_traffic {
    local iface="$1"

    echo "multigraph wireguard_peertraffic_$iface"

    if should_emit_config; then
        cat <<EOF
graph_title $iface peer traffic
graph_args --base 1000
graph_vlabel bits in (-) / out (+) per \${graph_period}
graph_category wireguard
graph_info This graph shows the traffic per peer on the $iface wireguard interface. Traffic is shown in bits per second.
EOF
    fi

    declare -a _peers
    wg_peers peers "$iface"

    for line in "${peers[@]}"; do
        IFS=$'\t' read -ra peer <<< "$line"
        peer_id="$(safe_peer_id "${peer[0]}")"

        if should_emit_config; then
            # List config for up/down values for each peer
            cat <<EOF
down_${peer_id}.label received
down_${peer_id}.type DERIVE
down_${peer_id}.graph no
down_${peer_id}.cdef down_${peer_id},8,*
down_${peer_id}.min 0
up_${peer_id}.label ${peer[2]}
up_${peer_id}.type DERIVE
up_${peer_id}.negative down_${peer_id}
up_${peer_id}.cdef up_${peer_id},8,*
up_${peer_id}.min 0
EOF
        fi

        if should_emit_values; then
            cat <<EOF
down_${peer_id}.value ${peer[5]}
up_${peer_id}.value ${peer[6]}
EOF
        fi
    done
}

case $1 in
    autoconf)
        if wg_exists; then
            echo "yes"
        else
            echo "no (wg command not found)"
        fi
        ;;
    suggest)
        if wg_exists; then
            declare -a ifaces
            wg_interfaces ifaces 1
            IFS=$'\n'
            echo "${ifaces[*]}"
        fi
        ;;
    *)
        peer_count

        declare -a ifaces
        wg_interfaces ifaces
        for iface in "${ifaces[@]}"; do
            peer_traffic "$iface"
        done
        ;;
esac