#!/bin/bash
set -uo pipefail

# ==============================
# Прогрев nft-сетов (запускать НА ШЛЮЗЕ)
# ==============================
# Заново резолвит все домены, которые dnsmasq помечает директивами nftset=,
# обращаясь напрямую к host-инстансу 127.0.0.1#5350. Прямой запрос обходит
# кэш pihole FTL: запрос, попавший в кэш FTL, до :5350 не доходит и элементы
# в inet filter bbrkn_v4/bbrkn_v6 не обновляет — они молча истекают по timeout
# и домен уходит мимо туннеля.
#
# Источник списка доменов — сам задеплоенный конфиг, поэтому отдельный список
# синхронизировать не нужно: скрипт всегда работает с тем, что реально
# загружен в dnsmasq после последнего успешного деплоя.
#
# Самовосстановление: dnsmasq добавляет IP в сет только при свежем ответе
# upstream и иногда "клинит" — основной процесс начинает отвечать на каждую
# попытку записи "nftset ... netlink: Error: cache initialization failed:
# Protocol error" и больше сеты не наполняет до рестарта. Если за прогон в
# журнале основного процесса появилась такая ошибка — рестартуем dnsmasq и
# прогреваем ещё раз.

# --- Конфигурация через переменные окружения ---
# Задеплоенный конфиг с директивами nftset= (источник списка доменов)
NFTSET_CONF="${NFTSET_CONF:-/etc/dnsmasq.d/90-nftset.conf}"
# Host-инстанс dnsmasq, который применяет nftset=
RESOLVER="${RESOLVER:-127.0.0.1}"
RESOLVER_PORT="${RESOLVER_PORT:-5350}"

# Состояние и лог. Общий каталог для cron и деплоя — lock у них один.
STATE_DIR="${STATE_DIR:-/var/lib/bbrkn}"
LOG="${LOG:-${STATE_DIR}/warm.log}"
LOCK="${LOCK:-${STATE_DIR}/warm.lock}"
LOG_MAX_BYTES="${LOG_MAX_BYTES:-4194304}"   # ротация лога, 4 МиБ
LOCK_WAIT="${LOCK_WAIT:-0}"   # секунд ждать чужой прогон (0 — сразу выйти)

# Параметры опроса
PARALLEL="${PARALLEL:-8}"          # одновременных запросов к :5350
DIG_TIMEOUT="${DIG_TIMEOUT:-3}"    # секунд на попытку
DIG_TRIES="${DIG_TRIES:-2}"

# Пинг выключен по умолчанию: пакеты, рождённые на самом шлюзе, идут через
# OUTPUT и не проходят цепочку prerouting_mangle — уходят в обычный WAN без
# метки и почти все отваливаются по таймауту. Включать только для отладки.
WARM_PING="${WARM_PING:-0}"
PING_TIMEOUT="${PING_TIMEOUT:-2}"

NFT="${NFT:-/usr/sbin/nft}"

# systemd-сервис host-инстанса dnsmasq (:5350) и самовосстановление
HOST_DNSMASQ_SVC="${HOST_DNSMASQ_SVC:-dnsmasq}"
SELF_HEAL="${SELF_HEAL:-1}"

# Под root (cron.d, CI) sudo не нужен
SUDO=""
[ "$(id -u)" -eq 0 ] || SUDO="sudo -n"

mkdir -p "$STATE_DIR"

# Один инстанс: медленный прогон не должен наложиться на следующий по cron.
exec 9>"$LOCK"
if ! flock -w "$LOCK_WAIT" 9; then
    echo "$(date -Is) SKIP предыдущий запуск ещё выполняется" >>"$LOG"
    exit 0
fi

log() { echo "$(date -Is) $*" >>"$LOG"; }

# Количество элементов в сете. У каждого элемента есть timeout, поэтому
# "expires" встречается ровно один раз на элемент. Ошибка здесь не фатальна —
# это только метрика прогресса.
set_count() {
    $SUDO "$NFT" list set inet filter "$1" 2>/dev/null | grep -c 'expires' || true
}

main_pid() {
    systemctl show -p MainPID --value "$HOST_DNSMASQ_SVC" 2>/dev/null
}

# Ошибки записи в nft-сет у основного процесса dnsmasq с момента <since>.
# Ошибки дочерних (TCP) процессов не считаем — они разовые; заклинивание
# видно именно по основному.
nftset_errors() {  # <pid> <since>
    $SUDO journalctl -q --no-pager _PID="$1" --since "$2" 2>/dev/null \
        | grep -c 'nftset.*Protocol error' || true
}

if [ ! -r "$NFTSET_CONF" ]; then
    log "FATAL не читается $NFTSET_CONF"
    echo "Ошибка: не читается $NFTSET_CONF (скрипт запускается на шлюзе)" >&2
    exit 1
fi

mapfile -t HOSTS < <(sed -n 's|^nftset=/\([^/]*\)/.*|\1|p' "$NFTSET_CONF" | sort -u)
if [ "${#HOSTS[@]}" -eq 0 ]; then
    log "FATAL в $NFTSET_CONF не найдено ни одной директивы nftset="
    echo "Ошибка: в $NFTSET_CONF не найдено ни одной директивы nftset=" >&2
    exit 1
fi

warm_one() {
    local host="$1" a aaaa
    a=$(dig +short +timeout="$DIG_TIMEOUT" +tries="$DIG_TRIES" \
        @"$RESOLVER" -p "$RESOLVER_PORT" A "$host" 2>/dev/null \
        | grep -E '^[0-9]+\.' | head -1)
    aaaa=$(dig +short +timeout="$DIG_TIMEOUT" +tries="$DIG_TRIES" \
        @"$RESOLVER" -p "$RESOLVER_PORT" AAAA "$host" 2>/dev/null \
        | grep -E '^[0-9a-fA-F]*:' | head -1)

    if [ -z "$a" ] && [ -z "$aaaa" ]; then
        echo "dnsfail $host"
        return
    fi

    # Пинг только подтверждает проходимость маршрута; в сет IP попал уже на
    # шаге резолва выше.
    if [ "$WARM_PING" = "1" ] && [ -n "$a" ]; then
        if ping -n -q -c1 -W "$PING_TIMEOUT" "$a" >/dev/null 2>&1; then
            echo "ok $host $a"
        else
            echo "noping $host $a"
        fi
    else
        echo "ok $host ${a:-$aaaa}"
    fi
}
export -f warm_one
export RESOLVER RESOLVER_PORT DIG_TIMEOUT DIG_TRIES WARM_PING PING_TIMEOUT

# Один прогон по всем доменам. Заполняет RESULTS и ERRORS (ошибки записи в
# сет у основного процесса dnsmasq за время прогона).
warm_all() {
    local since pid
    since=$(date '+%Y-%m-%d %H:%M:%S')
    pid=$(main_pid)
    RESULTS="$(printf '%s\n' "${HOSTS[@]}" \
        | xargs -r -P "$PARALLEL" -I{} bash -c 'warm_one "$@"' _ {})"
    sleep 1   # журнал пишется асинхронно
    ERRORS=0
    if [ -n "$pid" ] && [ "$pid" != "0" ]; then
        ERRORS=$(nftset_errors "$pid" "$since")
    fi
}

BEFORE_V4=$(set_count bbrkn_v4)
BEFORE_V6=$(set_count bbrkn_v6)
log "START domains=${#HOSTS[@]} conf=$NFTSET_CONF v4=$BEFORE_V4 v6=$BEFORE_V6"

warm_all
HEALED=0
if [ "$ERRORS" -gt 0 ] && [ "$SELF_HEAL" = "1" ]; then
    log "WEDGE $HOST_DNSMASQ_SVC pid=$(main_pid) nftset_errors=$ERRORS — рестарт и повторный прогрев"
    if $SUDO systemctl restart "$HOST_DNSMASQ_SVC"; then
        sleep 2
        HEALED=1
        warm_all
    else
        log "FATAL не удалось перезапустить $HOST_DNSMASQ_SVC"
    fi
fi

OK=$(grep -c '^ok '        <<<"$RESULTS" || true)
NOPING=$(grep -c '^noping ' <<<"$RESULTS" || true)
DNSFAIL=$(grep -c '^dnsfail ' <<<"$RESULTS" || true)

AFTER_V4=$(set_count bbrkn_v4)
AFTER_V6=$(set_count bbrkn_v6)
log "DONE ok=$OK noping=$NOPING dnsfail=$DNSFAIL nftset_errors=$ERRORS healed=$HEALED v4=${BEFORE_V4}->${AFTER_V4} v6=${BEFORE_V6}->${AFTER_V6}"
if [ "$DNSFAIL" -gt 0 ]; then
    log "dnsfail: $(grep '^dnsfail ' <<<"$RESULTS" | cut -d' ' -f2 | tr '\n' ' ')"
fi

echo "✔ Прогрето доменов: $OK, не резолвится: $DNSFAIL; bbrkn_v4 ${BEFORE_V4}->${AFTER_V4}, bbrkn_v6 ${BEFORE_V6}->${AFTER_V6}"
[ "$HEALED" = "1" ] && echo "⚠ dnsmasq заклинило на записи в nft-сет — перезапущен ($HOST_DNSMASQ_SVC)"

# Ротация лога
if [ -f "$LOG" ] && [ "$(stat -c%s "$LOG")" -gt "$LOG_MAX_BYTES" ]; then
    mv -f "$LOG" "${LOG}.1"
fi

# Не ноль, если сеты так и не наполняются — деплой покажет предупреждение.
if [ "$ERRORS" -gt 0 ]; then
    log "FAIL nftset_errors=$ERRORS после прогрева"
    echo "Ошибка: dnsmasq не пишет в nft-сеты (nftset_errors=$ERRORS), см. journalctl -u $HOST_DNSMASQ_SVC" >&2
    exit 1
fi
exit 0