#!/bin/sh
# wg0-watchdog - keep the WireGuard client tunnel up so SSH to the tunnel
# address works after boot, after roaming to a new network, and after the NAT
# mapping or endpoint IP goes stale.
#
# Decisions, oldest symptom first:
#   interface or peer gone            -> restart wg-quick@wg0.service
#   no PersistentKeepalive in kernel  -> set it (covers a config edited without reload)
#   handshake older than STALE        -> re-point the peer at the endpoint
#                                        (re-resolves the hostname, forces a handshake)
#   handshake older than STALE_RESTART-> restart wg-quick@wg0.service
#
# Restarts are gap-limited so a broken network cannot make the unit flap.
# Never removes peers, addresses, or routes it did not create.

set -eu

IFACE="${WG_WD_IFACE:-wg0}"
UNIT="${WG_WD_UNIT:-wg-quick@wg0.service}"
CONF="${WG_WD_CONF:-/etc/wireguard/$IFACE.conf}"
KEEPALIVE="${WG_WD_KEEPALIVE:-25}"
STALE="${WG_WD_STALE:-120}"              # seconds: poke the endpoint
STALE_RESTART="${WG_WD_STALE_RESTART:-300}"   # seconds: restart the unit
MIN_RESTART_GAP="${WG_WD_MIN_RESTART_GAP:-180}"  # seconds between restarts
STATE_DIR="${WG_WD_STATE_DIR:-/var/lib/wg0-watchdog}"

DRY=0
[ "${1:-}" = "--dry-run" ] && DRY=1

now=$(date +%s)
log() { printf 'wg0-watchdog: %s\n' "$*"; }
run() {
    if [ "$DRY" -eq 1 ]; then
        printf 'wg0-watchdog: DRY-RUN would run:'
        printf ' %s' "$@"
        printf '\n'
    else
        "$@"
    fi
}

read_state() { cat "$1" 2>/dev/null || echo 0; }
write_state() {
    if [ "$DRY" -eq 1 ]; then
        return 0
    fi
    mkdir -p "$STATE_DIR"
    printf '%s\n' "$2" > "$1"
}

restart_unit() {
    reason=$1
    last=$(read_state "$STATE_DIR/last-restart")
    case "$last" in ''|*[!0-9]*) last=0 ;; esac
    since=$((now - last))
    if [ "$since" -lt "$MIN_RESTART_GAP" ]; then
        log "restart of $UNIT suppressed ($reason); last restart ${since}s ago"
        return 0
    fi
    log "restarting $UNIT ($reason)"
    write_state "$STATE_DIR/last-restart" "$now"
    # --no-block: a wg-quick up stuck on endpoint resolution must not wedge the
    # watchdog; the next run reports whether the restart worked.
    run systemctl restart --no-block "$UNIT"
}

# --- 1. interface / peer present -------------------------------------------
if ! ip link show "$IFACE" >/dev/null 2>&1; then
    restart_unit "interface $IFACE is missing"
    exit 0
fi

peer=$(wg show "$IFACE" peers 2>/dev/null | head -n 1)
if [ -z "${peer:-}" ]; then
    restart_unit "no peer configured on $IFACE"
    exit 0
fi

# --- 2. persistent keepalive present ---------------------------------------
live_ka=$(wg show "$IFACE" persistent-keepalive 2>/dev/null | awk -v p="$peer" '$1 == p {print $2}')
[ -n "${live_ka:-}" ] || live_ka=0
if [ "$live_ka" != "$KEEPALIVE" ]; then
    log "persistent-keepalive is $live_ka, setting $KEEPALIVE"
    run wg set "$IFACE" peer "$peer" persistent-keepalive "$KEEPALIVE"
fi

# --- 3. handshake age -------------------------------------------------------
hs=$(wg show "$IFACE" latest-handshakes 2>/dev/null | awk -v p="$peer" '$1 == p {print $2}')
case "${hs:-}" in ''|*[!0-9]*) hs=0 ;; esac
if [ "$hs" -eq 0 ]; then
    age=999999
else
    age=$((now - hs))
fi

# endpoint currently in use by the kernel, and the cached last-known-good one
live_ep=$(wg show "$IFACE" endpoints 2>/dev/null | awk -v p="$peer" '$1 == p {print $2}')
cache_file="$STATE_DIR/endpoint"
cached_ep=""
[ -f "$cache_file" ] && cached_ep=$(cat "$cache_file" 2>/dev/null || true)

if [ "$age" -lt "$STALE" ]; then
    if [ -n "$live_ep" ]; then
        write_state "$cache_file" "$live_ep"
    fi
    log "healthy: handshake ${age}s ago, endpoint ${live_ep:-unknown}, keepalive ${KEEPALIVE}s"
    exit 0
fi

log "stale: last handshake ${age}s ago"

# Prefer the endpoint from the config file, freshly resolved. Fall back to the
# kernel's current endpoint, then to the cache - important because a DNS server
# from the tunnel config may be reachable only through this very tunnel.
target_ep=""
conf_ep=$(awk -F= 'tolower($1) ~ /^[[:space:]]*endpoint[[:space:]]*$/ { gsub(/[[:space:]]/, "", $2); print $2 }' "$CONF" 2>/dev/null | tail -n 1)
if [ -n "${conf_ep:-}" ]; then
    host=${conf_ep%:*}
    port=${conf_ep##*:}
    case "$host" in
        *[!0-9.]*)
            resolved=$(getent hosts "$host" 2>/dev/null | awk '{print $1; exit}' || true)
            ;;
        *)
            resolved=$host
            ;;
    esac
    if [ -n "${resolved:-}" ]; then
        target_ep="$resolved:$port"
    else
        log "could not resolve $host (tunnel DNS?), falling back"
    fi
fi
[ -n "$target_ep" ] || target_ep=${live_ep:-${cached_ep:-}}

if [ "$age" -ge "$STALE_RESTART" ]; then
    restart_unit "handshake ${age}s old"
    exit 0
fi

if [ -z "$target_ep" ]; then
    log "no endpoint to re-point at; will retry next run"
    exit 0
fi

if [ "$target_ep" = "$live_ep" ]; then
    log "endpoint already $target_ep; forcing a handshake"
fi
log "re-pointing peer at $target_ep"
run wg set "$IFACE" peer "$peer" endpoint "$target_ep"
exit 0