#!/usr/bin/env bash
# AVNI firewall: persistent blacklist (ipset+iptables, host + Docker) and a
# WHITELIST that can never be banned/blacklisted (fail2ban ignoreip + guard).
set -uo pipefail
SET=avni_blacklist
SAVE=/etc/avni/blacklist.save
WL=/etc/avni/whitelist.txt
F2B_CONF=/etc/fail2ban/jail.d/avni-whitelist.conf
mkdir -p /etc/avni; touch "$WL"

ensure_set(){ ipset list -n 2>/dev/null | grep -qx "$SET" || \
  ipset create "$SET" hash:net family inet hashsize 1024 maxelem 65536; }
ensure_rules(){
  iptables -C INPUT -m set --match-set "$SET" src -j DROP 2>/dev/null || \
    iptables -I INPUT 1 -m set --match-set "$SET" src -j DROP
  if iptables -L DOCKER-USER -n >/dev/null 2>&1; then
    iptables -C DOCKER-USER -m set --match-set "$SET" src -j DROP 2>/dev/null || \
      iptables -I DOCKER-USER 1 -m set --match-set "$SET" src -j DROP
  fi
}
save(){ ipset save "$SET" > "$SAVE" 2>/dev/null; }
my_ip(){ echo "${SSH_CLIENT%% *}"; }
is_wl(){ grep -qxF "$1" "$WL" 2>/dev/null; }
jails(){ fail2ban-client status 2>/dev/null | sed -n 's/.*Jail list:[[:space:]]*//p' | tr ',' ' '; }

gen_f2b(){   # write a persistent ignoreip drop-in from the whitelist + reload
  local ips; ips=$(tr '\n' ' ' < "$WL")
  { echo "# AUTO-GENERATED by avni-firewall.sh — do not edit; use: avni whitelist add/del"
    echo "[DEFAULT]"
    echo "ignoreip = 127.0.0.1/8 ::1 $ips"; } > "$F2B_CONF"
  fail2ban-client reload >/dev/null 2>&1 || true
}

resolve(){
  if [[ "$1" =~ ^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+(/[0-9]+)?$ ]]; then echo "$1"; return; fi
  getent ahostsv4 "$1" 2>/dev/null | awk '{print $1}' | sort -u
}

case "${1:-}" in
  init)    ensure_set; ensure_rules; save ;;
  restore) if [ -f "$SAVE" ]; then ipset restore -! < "$SAVE" 2>/dev/null; else ensure_set; fi; ensure_rules; gen_f2b ;;
  add)
    ensure_set; ensure_rules
    for ip in $(resolve "$2"); do
      if [ "$ip" = "$(my_ip)" ]; then echo "REFUSED: $ip is YOUR current SSH IP"; continue; fi
      if is_wl "$ip"; then echo "REFUSED: $ip is WHITELISTED (remove from whitelist first)"; continue; fi
      ipset add -! "$SET" "$ip" && echo "blacklisted $ip"
    done; save ;;
  del)
    for ip in $(resolve "$2"); do ipset del "$SET" "$ip" 2>/dev/null && echo "removed $ip" || echo "$ip not listed"; done; save ;;
  list)    ipset list "$SET" 2>/dev/null | sed -n '/Members:/,$p' | tail -n +2 ;;
  count)   ipset list "$SET" 2>/dev/null | sed -n '/Members:/,$p' | tail -n +2 | grep -c . ;;
  whitelist)
    case "${2:-list}" in
      add)
        for ip in $(resolve "$3"); do
          grep -qxF "$ip" "$WL" || echo "$ip" >> "$WL"
          ipset del "$SET" "$ip" 2>/dev/null    # un-blacklist if present
          for j in $(jails); do
            fail2ban-client set "$j" addignoreip "$ip" >/dev/null 2>&1
            fail2ban-client set "$j" unbanip "$ip"   >/dev/null 2>&1
          done
          echo "whitelisted + unbanned $ip (can never be banned/blacklisted)"
        done
        sort -u "$WL" -o "$WL"; gen_f2b; save ;;
      del)
        for ip in $(resolve "$3"); do
          grep -vxF "$ip" "$WL" > "$WL.tmp" && mv "$WL.tmp" "$WL"
          for j in $(jails); do fail2ban-client set "$j" delignoreip "$ip" >/dev/null 2>&1; done
          echo "removed $ip from whitelist"
        done; gen_f2b ;;
      list|*) [ -s "$WL" ] && cat "$WL" || echo "(whitelist empty)" ;;
    esac ;;
  *) echo "usage: $0 {init|restore|add <ip|cidr|domain>|del <ip>|list|count|whitelist {add|del|list} <ip>}"; exit 1 ;;
esac