#!/bin/sh
# =============================================================================
# Data-Shield IPv4 Blocklist - Script de mise à jour pour FreeBSD / PF
# Version : 1.8.0
# Usage   : /usr/local/sbin/update_datashield.sh [--cron]
# Cron    : 0 */6 * * * root /usr/local/sbin/update_datashield.sh --cron
# Dépendance : ipcalc  →  pkg install ipcalc
# =============================================================================

# ---------------------------------------------------------------------------
# CONFIGURATION
# ---------------------------------------------------------------------------

# Miroirs source — tentés dans l'ordre jusqu'au premier succès (liste complète ~100k IPs)
BLOCKLIST_URLS="
https://raw.githubusercontent.com/duggytuxy/Data-Shield_IPv4_Blocklist/refs/heads/main/prod_data-shield_ipv4_blocklist.txt
https://gitlab.com/duggytuxy/Data-Shield-IPv4-Blocklist/-/raw/main/prod_data-shield_ipv4_blocklist.txt?ref_type=heads
https://cdn.jsdelivr.net/gh/duggytuxy/Data-Shield_IPv4_Blocklist@refs/heads/main/prod_data-shield_ipv4_blocklist.txt
https://bitbucket.org/duggytuxy/data-shield-ipv4-blocklist/raw/HEAD/prod_data-shield_ipv4_blocklist.txt
https://codeberg.org/duggytuxy21/Data-Shield_IPv4_Blocklist/raw/branch/main/prod_data-shield_ipv4_blocklist.txt
"

# Répertoire de travail
BLOCKLIST_DIR="/etc/pf"
BLOCKLIST_FILE="${BLOCKLIST_DIR}/datashield_blocklist.txt"
BLOCKLIST_TMP="${BLOCKLIST_DIR}/datashield_blocklist.tmp"
BLOCKLIST_BACKUP="${BLOCKLIST_DIR}/datashield_blocklist.bak"

# Fichier whitelist — une entrée par ligne, commentaires # supportés
# Supporte les IPs individuelles et les subnets CIDR
WHITELIST_FILE="${BLOCKLIST_DIR}/datashield_whitelist.txt"

# Logs
LOG_FILE="/var/log/datashield_update.log"
LOG_MAX_LINES=500

# Seuils de validation
MIN_IPS=10000
MAX_IPS=200000

# Nom de la table PF
PF_TABLE="datashield"

# Timeout fetch (secondes)
FETCH_TIMEOUT=60

# Préfixe CIDR minimum accepté en whitelist (ex: 16 = /16 max, refus de /8, /0, etc.)
CIDR_MIN_PREFIX=16

# ---------------------------------------------------------------------------
# FONCTIONS UTILITAIRES
# ---------------------------------------------------------------------------

timestamp() { date "+%Y-%m-%dT%H:%M:%S"; }

log() {
    LEVEL="$1"; shift
    MSG="[$(timestamp)] [$LEVEL] $*"
    if [ "$CRON_MODE" = "1" ]; then
        # Mode cron : uniquement WARN et ERROR pour éviter les mails inutiles
        case "$LEVEL" in
            "WARN "|"ERROR") echo "$MSG" ;;
        esac
    else
        # Mode interactif : tout afficher
        echo "$MSG"
    fi
    echo "$MSG" >> "$LOG_FILE"
}

log_info()  { log "INFO " "$@"; }
log_warn()  { log "WARN " "$@"; }
log_error() { log "ERROR" "$@"; }
log_ok()    { log "OK   " "$@"; }

rotate_log() {
    if [ -f "$LOG_FILE" ]; then
        LINES=$(wc -l < "$LOG_FILE")
        if [ "$LINES" -gt "$LOG_MAX_LINES" ]; then
            mv "$LOG_FILE" "${LOG_FILE}.old"
            log_info "Log rotaté vers ${LOG_FILE}.old (${LINES} lignes)"
        fi
    fi
}

cleanup() { rm -f "$BLOCKLIST_TMP" "${BLOCKLIST_TMP}.wl" "${BLOCKLIST_TMP}.patterns"; }

die() {
    log_error "$@"
    cleanup
    exit 1
}

# ---------------------------------------------------------------------------
# VÉRIFICATIONS PRÉALABLES
# ---------------------------------------------------------------------------

preflight_checks() {
    log_info "=== Démarrage de la mise à jour Data-Shield ==="

    [ "$(id -u)" -ne 0 ] && die "Ce script doit être exécuté en root"

    command -v pfctl > /dev/null 2>&1 || die "pfctl introuvable — PF est-il activé ?"

    pfctl -si > /dev/null 2>&1 || die "PF ne répond pas — vérifiez pf_enable=\"YES\" dans /etc/rc.conf"

    command -v ipcalc > /dev/null 2>&1 || die "ipcalc introuvable — lancez : pkg install ipcalc"

    if [ ! -d "$BLOCKLIST_DIR" ]; then
        mkdir -p "$BLOCKLIST_DIR" || die "Impossible de créer $BLOCKLIST_DIR"
        log_info "Répertoire $BLOCKLIST_DIR créé"
    fi

    log_ok "Vérifications préalables OK"
}

# ---------------------------------------------------------------------------
# TÉLÉCHARGEMENT AVEC FALLBACK MULTI-MIROIRS
# ---------------------------------------------------------------------------

download_blocklist() {
    DOWNLOAD_OK=0

    for URL in $BLOCKLIST_URLS; do
        # Ignorer les lignes vides (découpage du heredoc)
        [ -z "$URL" ] && continue

        log_info "Tentative de téléchargement depuis : $URL"

        # Syntaxe fetch(1) FreeBSD :
        #   -T  timeout en secondes
        #   -q  silencieux
        #   -o  fichier de sortie
        #   --no-verify-peer  flag booléen (sans argument)
        fetch -T "$FETCH_TIMEOUT" -q --no-verify-peer -o "$BLOCKLIST_TMP" "$URL"
        FETCH_EXIT=$?

        if [ $FETCH_EXIT -ne 0 ]; then
            log_warn "Miroir indisponible (code $FETCH_EXIT) : $URL"
            rm -f "$BLOCKLIST_TMP"
            continue
        fi

        if [ ! -s "$BLOCKLIST_TMP" ]; then
            log_warn "Fichier vide reçu depuis : $URL"
            rm -f "$BLOCKLIST_TMP"
            continue
        fi

        log_ok "Téléchargement réussi depuis : $URL ($(wc -c < "$BLOCKLIST_TMP" | tr -d ' ') octets)"
        DOWNLOAD_OK=1
        break
    done

    [ "$DOWNLOAD_OK" -eq 0 ] && die "Tous les miroirs ont échoué — abandon"
}

# ---------------------------------------------------------------------------
# VALIDATION DU CONTENU
# ---------------------------------------------------------------------------

validate_blocklist() {
    log_info "Validation du contenu téléchargé..."

    IP_COUNT=$(grep -E '^([0-9]{1,3}\.){3}[0-9]{1,3}$' "$BLOCKLIST_TMP" 2>/dev/null | wc -l | tr -d ' ')

    log_info "IPs valides détectées : $IP_COUNT"

    [ "$IP_COUNT" -lt "$MIN_IPS" ] && \
        die "Liste suspecte : seulement $IP_COUNT IPs (minimum : $MIN_IPS) — abandon"

    [ "$IP_COUNT" -gt "$MAX_IPS" ] && \
        die "Liste suspecte : $IP_COUNT IPs dépassent le maximum ($MAX_IPS) — abandon"

    SUSPICIOUS=$(grep -vE '^(#.*)?$|^([0-9]{1,3}\.){3}[0-9]{1,3}$' "$BLOCKLIST_TMP" 2>/dev/null | wc -l | tr -d ' ')
    [ "${SUSPICIOUS:-0}" -gt 100 ] && \
        log_warn "Lignes non-IP détectées : $SUSPICIOUS — vérification recommandée"

    log_ok "Validation OK : $IP_COUNT entrées"
}

# ---------------------------------------------------------------------------
# WHITELIST — GESTION IP SIMPLE ET SUBNET CIDR
# ---------------------------------------------------------------------------

# Convertit une IP en entier 32 bits
ip_to_int() {
    IFS='.' read -r A B C D << EOF
$1
EOF
    echo $(( (A << 24) | (B << 16) | (C << 8) | D ))
}

# Convertit un entier 32 bits en notation IPv4
int_to_ip() {
    N="$1"
    echo "$(( (N >> 24) & 0xFF )).$(( (N >> 16) & 0xFF )).$(( (N >> 8) & 0xFF )).$(( N & 0xFF ))"
}

# Génère toutes les IPs d'un CIDR dans le fichier $2
# Usage : expand_cidr "192.168.1.0/24" "/tmp/patterns"
expand_cidr() {
    CIDR="$1"
    OUT="$2"

    NETWORK=$(echo "$CIDR" | cut -d/ -f1)
    PREFIX=$(echo "$CIDR" | cut -d/ -f2)

    if ! echo "$PREFIX" | grep -qE '^[0-9]+$' || [ "$PREFIX" -lt 0 ] || [ "$PREFIX" -gt 32 ]; then
        return 1
    fi

    if [ "$PREFIX" -lt "$CIDR_MIN_PREFIX" ]; then
        log_warn "CIDR /$PREFIX refusé : préfixe inférieur au minimum /$CIDR_MIN_PREFIX (modifier CIDR_MIN_PREFIX si nécessaire)"
        return 1
    fi

    NET_INT=$(ip_to_int "$NETWORK")

    if [ "$PREFIX" -eq 32 ]; then
        int_to_ip "$NET_INT" >> "$OUT"
        return 0
    fi

    HOST_BITS=$(( 32 - PREFIX ))
    # SIZE = 2^HOST_BITS via décalage shell
    SIZE=$(( 1 << HOST_BITS ))
    MASK=$(( 0xFFFFFFFF ^ (SIZE - 1) ))
    FIRST=$(( NET_INT & MASK ))
    LAST=$(( FIRST + SIZE - 1 ))

    I=$FIRST
    while [ "$I" -le "$LAST" ]; do
        int_to_ip "$I" >> "$OUT"
        I=$(( I + 1 ))
    done
}

apply_whitelist() {
    if [ ! -f "$WHITELIST_FILE" ]; then
        log_info "Fichier whitelist absent ($WHITELIST_FILE) — étape ignorée"
        return 0
    fi

    ENTRY_COUNT=$(grep -cvE '^\s*(#.*)?$' "$WHITELIST_FILE" 2>/dev/null || echo 0)

    if [ "$ENTRY_COUNT" -eq 0 ]; then
        log_info "Fichier whitelist vide — étape ignorée"
        return 0
    fi

    log_info "Application de la whitelist : $WHITELIST_FILE ($ENTRY_COUNT entrées)"

    # Fichier temporaire accumulant tous les patterns à exclure (IPs et plages CIDR)
    WL_PATTERNS="${BLOCKLIST_TMP}.patterns"
    rm -f "$WL_PATTERNS"

    while IFS= read -r RAW_ENTRY; do

        # Ignorer commentaires et lignes vides
        echo "$RAW_ENTRY" | grep -qE '^\s*(#.*)?$' && continue

        # Nettoyer les espaces et tabulations
        ENTRY=$(echo "$RAW_ENTRY" | tr -d ' \t')
        [ -z "$ENTRY" ] && continue

        # ---- CAS 1 : SUBNET CIDR ----------------------------------------
        if echo "$ENTRY" | grep -qE '^([0-9]{1,3}\.){3}[0-9]{1,3}/[0-9]{1,2}$'; then

            if ! ipcalc "$ENTRY" > /dev/null 2>&1; then
                log_warn "CIDR invalide ignoré : $ENTRY"
                continue
            fi

            PREFIX=$(echo "$ENTRY" | cut -d/ -f2)
            HOST_BITS=$(( 32 - PREFIX ))
            SIZE=$(( 1 << HOST_BITS ))
            log_info "Expansion du subnet : $ENTRY ($SIZE adresses)"
            expand_cidr "$ENTRY" "$WL_PATTERNS"

        # ---- CAS 2 : IP INDIVIDUELLE -------------------------------------
        elif echo "$ENTRY" | grep -qE '^([0-9]{1,3}\.){3}[0-9]{1,3}$'; then
            echo "$ENTRY" >> "$WL_PATTERNS"

        # ---- CAS 3 : FORMAT INVALIDE -------------------------------------
        else
            log_warn "Entrée whitelist ignorée (format invalide) : $ENTRY"
        fi

    done < "$WHITELIST_FILE"

    # Aucun pattern généré
    if [ ! -s "$WL_PATTERNS" ]; then
        log_info "Whitelist appliquée : aucun pattern généré"
        rm -f "$WL_PATTERNS"
        return 0
    fi

    PATTERN_COUNT=$(wc -l < "$WL_PATTERNS" | tr -d ' ')
    log_info "Filtrage en une passe : $PATTERN_COUNT patterns à exclure"

    BEFORE=$(grep -cE '^([0-9]{1,3}\.){3}[0-9]{1,3}$' "$BLOCKLIST_TMP" 2>/dev/null || echo 0)

    # Un seul grep -vFf pour toute la whitelist
    grep -vFf "$WL_PATTERNS" "$BLOCKLIST_TMP" > "${BLOCKLIST_TMP}.wl" \
        && mv "${BLOCKLIST_TMP}.wl" "$BLOCKLIST_TMP"

    rm -f "$WL_PATTERNS"

    AFTER=$(grep -cE '^([0-9]{1,3}\.){3}[0-9]{1,3}$' "$BLOCKLIST_TMP" 2>/dev/null || echo 0)
    TOTAL_REMOVED=$(( BEFORE - AFTER ))

    if [ "$TOTAL_REMOVED" -gt 0 ]; then
        log_ok "Whitelist appliquée : $TOTAL_REMOVED entrée(s) retirée(s) de la blacklist"
    else
        log_info "Whitelist appliquée : aucune correspondance trouvée"
    fi
}

# ---------------------------------------------------------------------------
# COMPARAISON AVEC L'ANCIENNE LISTE
# ---------------------------------------------------------------------------

compare_with_previous() {
    if [ ! -f "$BLOCKLIST_FILE" ]; then
        log_info "Première installation — pas de comparaison possible"
        return 0
    fi

    OLD_COUNT=$(grep -E '^([0-9]{1,3}\.){3}[0-9]{1,3}$' "$BLOCKLIST_FILE" 2>/dev/null | wc -l | tr -d ' ')
    NEW_COUNT=$(grep -E '^([0-9]{1,3}\.){3}[0-9]{1,3}$' "$BLOCKLIST_TMP"  2>/dev/null | wc -l | tr -d ' ')
    DIFF=$((NEW_COUNT - OLD_COUNT))

    if [ "$DIFF" -ge 0 ]; then
        log_info "Évolution : +${DIFF} IPs (${OLD_COUNT} → ${NEW_COUNT})"
    else
        log_info "Évolution : ${DIFF} IPs (${OLD_COUNT} → ${NEW_COUNT})"
    fi

    if [ "$OLD_COUNT" -gt 0 ]; then
        THRESHOLD=$((OLD_COUNT / 2))
        if [ "$NEW_COUNT" -lt "$THRESHOLD" ]; then
            log_warn "ATTENTION : nouvelle liste > 50% plus petite que l'ancienne — vérification recommandée"
        fi
    fi
}

# ---------------------------------------------------------------------------
# REMPLACEMENT ATOMIQUE + BACKUP
# ---------------------------------------------------------------------------

install_blocklist() {
    if [ -f "$BLOCKLIST_FILE" ]; then
        cp "$BLOCKLIST_FILE" "$BLOCKLIST_BACKUP"
        log_info "Backup : $BLOCKLIST_BACKUP"
    fi

    mv "$BLOCKLIST_TMP" "$BLOCKLIST_FILE" || die "Impossible d'installer la nouvelle liste"
    chmod 600 "$BLOCKLIST_FILE"

    log_ok "Nouvelle liste installée : $BLOCKLIST_FILE"
}

# ---------------------------------------------------------------------------
# RECHARGEMENT DE LA TABLE PF
# ---------------------------------------------------------------------------

reload_pf_table() {
    log_info "Rechargement de la table PF <${PF_TABLE}>..."
    if ! pfctl -t "$PF_TABLE" -T show > /dev/null 2>&1; then
        log_warn "Table <${PF_TABLE}> absente de PF"
        log_warn "Vérifiez que /etc/pf.conf contient : table <${PF_TABLE}> persist file \"${BLOCKLIST_FILE}\""
        return 0
    fi
    
    pfctl -t "$PF_TABLE" -T replace -f "$BLOCKLIST_FILE" > /dev/null 2>&1
    PF_EXIT=$?
    
    if [ $PF_EXIT -ne 0 ]; then
        log_error "Échec du rechargement PF (code $PF_EXIT)"
        if [ -f "$BLOCKLIST_BACKUP" ]; then
            log_warn "Restauration du backup en cours..."
            pfctl -t "$PF_TABLE" -T replace -f "$BLOCKLIST_BACKUP" \
                && log_ok "Backup restauré avec succès" \
                || log_error "Restauration échouée — intervention manuelle requise"
        fi
        die "Rechargement PF échoué"
    fi
    
    LOADED=$(pfctl -t "$PF_TABLE" -T show 2>/dev/null | wc -l | tr -d ' ')
    log_ok "Table PF <${PF_TABLE}> active : ${LOADED} entrées"
}

# ---------------------------------------------------------------------------
# RAPPORT FINAL
# ---------------------------------------------------------------------------

final_report() {
    FINAL_COUNT=$(grep -E '^([0-9]{1,3}\.){3}[0-9]{1,3}$' "$BLOCKLIST_FILE" 2>/dev/null | wc -l | tr -d ' ')
    log_ok "=== Mise à jour terminée : ${FINAL_COUNT} IPs en protection ==="
}

# ---------------------------------------------------------------------------
# POINT D'ENTRÉE
# ---------------------------------------------------------------------------

main() {
    rotate_log
    preflight_checks
    download_blocklist
    validate_blocklist
    apply_whitelist
    compare_with_previous
    install_blocklist
    reload_pf_table
    final_report
}

trap 'log_warn "Script interrompu (signal reçu)"; cleanup; exit 130' INT TERM

# Parsing des arguments
CRON_MODE=0
for ARG in "$@"; do
    case "$ARG" in
        --cron) CRON_MODE=1 ;;
        *) echo "Usage: $0 [--cron]" >&2; exit 1 ;;
    esac
done

main
exit 0