#!/bin/bash
# Author: Christophe Casalegno / Brain 0verride
# Contact: christophe.casalegno@scalarx.com
# sxwall / ScalarX Wall
# Version 1.7
#
# Copyright (c) 2020 Christophe Casalegno
#
# This program is free software: you can redistribute it and/or modify
#
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program.  If not, see <https://www.gnu.org/licenses/>
#
# The license is available on this server here:
# https://www.christophe-casalegno.com/licences/gpl-3.0.txt

function root_check()
{
    if [[ $EUID -ne 0 ]]
    then
        echo "Error: This script must be run as root." >&2
        exit 1
    fi
}

function lock_check()
{
    command -v flock >/dev/null 2>&1 || { echo "Error: flock not found." >&2; exit 1; }

    local LOCKDIR="/run/lock"
    mkdir -p "$LOCKDIR" 2>/dev/null || true

    local NAME
    NAME="${0##*/}"
    local LOCKFILE="$LOCKDIR/${NAME}.lock"

    if [[ -L "$LOCKFILE" ]]
    then
        echo "Error: lockfile is a symlink ($LOCKFILE)" >&2
        exit 1
    fi

    if [[ -e "$LOCKFILE" && ! -f "$LOCKFILE" ]]
    then
        echo "Error: lockfile is not a regular file ($LOCKFILE)" >&2
        exit 1
    fi

    if [[ -e "$LOCKFILE" ]]
    then
        local OWNER_UID
        OWNER_UID="$(stat -c %u "$LOCKFILE" 2>/dev/null)" || { echo "Error: Cannot stat $LOCKFILE" >&2; exit 1; }
        if [[ "$OWNER_UID" != "0" ]]
        then
            echo "Error: lockfile owner is not root ($LOCKFILE, uid=$OWNER_UID)" >&2
            exit 1
        fi
    fi

    local OLDUMASK
    OLDUMASK="$(umask)"
    umask 077

    local LOCKFD
    exec {LOCKFD}<>"$LOCKFILE" || { echo "Error: Cannot open lock $LOCKFILE" >&2; umask "$OLDUMASK"; exit 1; }

    if ! flock -n "$LOCKFD"
    then
        local PID
        PID="$(cat "$LOCKFILE" 2>/dev/null)"
        echo "Error: $NAME is already running (PID ${PID:-unknown})." >&2
        umask "$OLDUMASK"
        exit 1
    fi

    if [[ -w "/proc/self/fd/$LOCKFD" ]]
    then
        : > "/proc/self/fd/$LOCKFD"
    else
        : > "$LOCKFILE"
    fi

    printf '%s\n' "$$" >&"$LOCKFD"

    umask "$OLDUMASK"
}

root_check
lock_check

SCRIPT_NAME="${0##*/}"
WORKPLACE='/home/sxadmin/data/sxwall'
DATADIR="$WORKPLACE/data"
BACKUPDIR="$WORKPLACE/backups"
IPTABLES_FILE='/etc/iptables.rules'

# Aggregated ipblocks URL 
URL_V4='https://www.ipdeny.com/ipblocks/data/aggregated'
URL_V6='https://www.ipdeny.com/ipv6/ipaddresses/aggregated'

# set names 
WHITE_V4='sx_white_v4' ; WHITE_V6='sx_white_v6'
MASTER_V4="sx_master_v4"  ; MASTER_V6="sx_master_v6"
MANUAL_V4="sx_manual_v4"  ; MANUAL_V6="sx_manual_v6"

# Whitelist
WHITELIST=("1.1.1.1" "8.8.8.8" "9.9.9.9" "127.0.0.1" "::1" "192.168.1.0/24")
WHITELIST_DIR_DEFAULT="/etc/sxwall/whitelist.d"
WHITELIST_DIR="${WHITELIST_DIR:-$WHITELIST_DIR_DEFAULT}"
WHITELIST_GLOB="${WHITELIST_GLOB:-*.conf}"

# Predefined blacklists
PREDEF_V4="sx_predef_v4"
PREDEF_V6="sx_predef_v6"

BLACKLIST_DIR_DEFAULT="/etc/sxwall/blacklist.d"
BLACKLIST_DIR="${BLACKLIST_DIR:-$BLACKLIST_DIR_DEFAULT}"
BLACKLIST_GLOB="${BLACKLIST_GLOB:-*.conf}"


# Zones
# AF: Africa
# AS: Asia
# EU: Europe
# NA: North America
# SA: South America
# OC: Oceania
# AN: Antartic

AF='dz ao bj bw bf bi cm cv cf td km cg cd ci dj eg gq er et ga gm gh gn gw ke ls lr ly mg mw ml mr mu yt ma mz na ne ng re rw st sn sc sl so za ss sd sz tz tg tn ug zm zw'
AS='af am az bh bd bt bn kh cn cc ge hk in id ir iq il jp jo kz kp kr kw kg la lb mo my mv mn mm np om pk ph qa sa sg lk sy tw tj th tr tm ae uz vn ye'
EU='al ad at by be ba bg hr cy cz dk ee fo fi fr de gi gr hu is ie it lv li lt lu mk mt md mc me nl no pl pt ro ru sm rs sk si es se ch ua gb va'
NA='ag bs bb bz bm vg ca ky cr cu dm do sv gl gd gp gt ht hn jm mq mx ms ni pa pr bl kn lc mf pm vc tt tc us vi'
SA='ar bo br cl co ec fk gf gy py pe sr uy ve'
OC='as au nz ck fj pf gu ki mh fm nr nc nu nf mp pw pg ws sb tk to tv vu wf'
AN='aq'

ZONES_TO_UPDATE="$AF $AS $EU $NA $SA $OC $AN"

ACTION="$1" ; TARGET="$2" ; PARAM="$3"

function region_to_countries()
{
    local r="$1"
    case "$r" in
        AF) echo "$AF" ;;
        AS) echo "$AS" ;;
        EU) echo "$EU" ;;
        NA) echo "$NA" ;;
        SA) echo "$SA" ;;
        OC) echo "$OC" ;;
        AN) echo "$AN" ;;
        *)  echo "" ;;
    esac
}

# Help me! 
#
function show_usage() 
{
    echo "SxWall - Usage:"
    echo "$SCRIPT_NAME start                     : Initialize firewall"
    echo "$SCRIPT_NAME reload			 : Reload whitelist/blacklist and reinject rules"
    echo "$SCRIPT_NAME check {target}            : Check status (IP/CIDR/Country/ASN/Region)"
    echo "$SCRIPT_NAME clean                     : Flush dynamic/manual blocks (keeps whitelist + predef blacklist)"
    echo "$SCRIPT_NAME import {file|-} [--clear] : Import IP/CIDR list into manual sets"
    echo "$SCRIPT_NAME list                      : Show active blocks"
    echo "$SCRIPT_NAME update                    : Update zones & ASNs"
    echo "$SCRIPT_NAME update asn {AS}           : Force update ASN"
    echo "$SCRIPT_NAME save                      : Backup sets"
    echo "$SCRIPT_NAME restore {date|latest}     : Restore backup"
    echo "$SCRIPT_NAME whoas {IP}                : Identify ASN"
    echo "$SCRIPT_NAME block {target}            : Block IP, CIDR, Country, Zone or ASN"
    echo "$SCRIPT_NAME unblock {target}          : Unblock target"
    echo ''
    echo 'Zones: AF: Africa, AS: Asia, EU: Europe, NA: North America, SA: South America, OC: Oceania, AN: Antartic'

    exit 1
}

function _is_ipv4()
{
    local ip="$1" o1 o2 o3 o4 x

    IFS='.' read -r o1 o2 o3 o4 <<< "$ip"

    [ -n "$o1" ] && [ -n "$o2" ] && [ -n "$o3" ] && [ -n "$o4" ] || return 1

    for x in "$o1" "$o2" "$o3" "$o4"
    do
        case "$x" in
            ''|*[!0-9]*)
                return 1
            ;;
        esac
        [ "$x" -ge 0 ] 2>/dev/null && [ "$x" -le 255 ] 2>/dev/null || return 1
    done

    return 0
}

function _is_ipv6()
{
    case "$1" in
        */*) return 1 ;;
        *:*) return 0 ;;
        *)   return 1 ;;
    esac
}

function _is_cidr_v4()
{
    local s="$1" ip pfx
    case "$s" in
        */*)
            ip="${s%/*}"
            pfx="${s#*/}"
        ;;
        *)
            return 1
        ;;
    esac

    _is_ipv4 "$ip" || return 1

    case "$pfx" in
        ''|*[!0-9]*)
            return 1
        ;;
    esac

    [ "$pfx" -ge 0 ] 2>/dev/null && [ "$pfx" -le 32 ] 2>/dev/null || return 1
    return 0
}

function _is_cidr_v6()
{
    local s="$1" ip pfx
    case "$s" in
        */*)
            ip="${s%/*}"
            pfx="${s#*/}"
        ;;
        *)
            return 1
        ;;
    esac

    # Cheap check: must look like v6 at least
    case "$ip" in
        *:*)
            ;;
        *)
            return 1
        ;;
    esac

    case "$pfx" in
        ''|*[!0-9]*)
            return 1
        ;;
    esac

    [ "$pfx" -ge 0 ] 2>/dev/null && [ "$pfx" -le 128 ] 2>/dev/null || return 1
    return 0
}

function _is_ip_or_cidr()
{
    _is_cidr_v4 "$1" && return 0
    _is_cidr_v6 "$1" && return 0
    _is_ipv4 "$1" && return 0
    _is_ipv6 "$1" && return 0
    return 1
}

function _is_symlink()
{
    [ -L "$1" ]
}

function blacklist_dir_check()
{
    local d="$1"

    [ -n "$d" ] || return 0
    [ -e "$d" ] || return 0

    if _is_symlink "$d"
    then
        echo "Error: BLACKLIST_DIR is a symlink ($d)" >&2
        exit 1
    fi

    if [ ! -d "$d" ]
    then
        echo "Error: BLACKLIST_DIR is not a directory ($d)" >&2
        exit 1
    fi

    local uid mode
    uid="$(stat -c %u "$d" 2>/dev/null)" || { echo "Error: Cannot stat $d" >&2; exit 1; }
    mode="$(stat -c %a "$d" 2>/dev/null)" || { echo "Error: Cannot stat $d" >&2; exit 1; }

    if [ "$uid" != "0" ]
    then
        echo "Error: BLACKLIST_DIR owner is not root ($d, uid=$uid)" >&2
        exit 1
    fi

    case "$mode" in
        ???)
            local g="${mode#?}"; g="${g%?}"
            local o="${mode#??}"
            case "$g" in 2|3|6|7) echo "Error: BLACKLIST_DIR is group-writable ($d, mode=$mode)" >&2; exit 1 ;; esac
            case "$o" in 2|3|6|7) echo "Error: BLACKLIST_DIR is world-writable ($d, mode=$mode)" >&2; exit 1 ;; esac
        ;;
        *)
            echo "Error: BLACKLIST_DIR mode unreadable ($d, mode=$mode)" >&2
            exit 1
        ;;
    esac
}

function blacklist_file_check()
{
    local f="$1"

    if _is_symlink "$f"
    then
        echo "Error: Blacklist file is a symlink ($f)" >&2
        exit 1
    fi

    if [ ! -f "$f" ]
    then
        echo "Error: Blacklist file is not a regular file ($f)" >&2
        exit 1
    fi

    local uid mode
    uid="$(stat -c %u "$f" 2>/dev/null)" || { echo "Error: Cannot stat $f" >&2; exit 1; }
    mode="$(stat -c %a "$f" 2>/dev/null)" || { echo "Error: Cannot stat $f" >&2; exit 1; }

    if [ "$uid" != "0" ]
    then
        echo "Error: Blacklist file owner is not root ($f, uid=$uid)" >&2
        exit 1
    fi

    case "$mode" in
        ???)
            local g="${mode#?}"; g="${g%?}"
            local o="${mode#??}"
            case "$g" in 2|3|6|7) echo "Error: Blacklist file is group-writable ($f, mode=$mode)" >&2; exit 1 ;; esac
            case "$o" in 2|3|6|7) echo "Error: Blacklist file is world-writable ($f, mode=$mode)" >&2; exit 1 ;; esac
        ;;
        *)
            echo "Error: Blacklist file mode unreadable ($f, mode=$mode)" >&2
            exit 1
        ;;
    esac
}

function load_blacklist_file()
{
    local f="$1"
    blacklist_file_check "$f"

    local ip

    while IFS= read -r ip
    do
        [ -n "$ip" ] || continue

        if ! _is_ip_or_cidr "$ip"
        then
            echo "Error: invalid IP/CIDR in blacklist file $f: $ip" >&2
            exit 1
        fi

        case "$ip" in
            *:*)
                ipset add "$PREDEF_V6" "$ip" -exist 2>/dev/null || { echo "Error: failed to add to $PREDEF_V6: $ip" >&2; exit 1; }
            ;;
            *)
                ipset add "$PREDEF_V4" "$ip" -exist 2>/dev/null || { echo "Error: failed to add to $PREDEF_V4: $ip" >&2; exit 1; }
            ;;
        esac
    done < <(awk '
        /^[[:space:]]*$/ { next }
        /^[[:space:]]*#/ { next }
        {
            sub(/[[:space:]]*#.*/, "", $0)
            gsub(/[[:space:]]+/, "", $0)
            if ($0 != "") print $0
        }
    ' "$f")
}

function load_blacklist_dir()
{
    local d="$1"
    [ -n "$d" ] || return 0
    [ -d "$d" ] || return 0

    blacklist_dir_check "$d"

    shopt -s nullglob
    local f
    for f in "$d"/$BLACKLIST_GLOB
    do
        load_blacklist_file "$f"
    done
    shopt -u nullglob
}

function whitelist_dir_check()
{
    local d="$1"

    [ -n "$d" ] || return 0
    [ -e "$d" ] || return 0

    if _is_symlink "$d"
    then
        echo "Error: WHITELIST_DIR is a symlink ($d)" >&2
        exit 1
    fi

    if [ ! -d "$d" ]
    then
        echo "Error: WHITELIST_DIR is not a directory ($d)" >&2
        exit 1
    fi

    local uid mode
    uid="$(stat -c %u "$d" 2>/dev/null)" || { echo "Error: Cannot stat $d" >&2; exit 1; }
    mode="$(stat -c %a "$d" 2>/dev/null)" || { echo "Error: Cannot stat $d" >&2; exit 1; }

    if [ "$uid" != "0" ]
    then
        echo "Error: WHITELIST_DIR owner is not root ($d, uid=$uid)" >&2
        exit 1
    fi

    # Reject group/world writable dirs (octal last 2 digits)
    case "$mode" in
        ???)
            local g="${mode#?}"; g="${g%?}"
            local o="${mode#??}"
            case "$g" in 2|3|6|7) echo "Error: WHITELIST_DIR is group-writable ($d, mode=$mode)" >&2; exit 1 ;; esac
            case "$o" in 2|3|6|7) echo "Error: WHITELIST_DIR is world-writable ($d, mode=$mode)" >&2; exit 1 ;; esac
        ;;
        *)
            # If mode is weird, be strict
            echo "Error: WHITELIST_DIR mode unreadable ($d, mode=$mode)" >&2
            exit 1
        ;;
    esac
}

function whitelist_file_check()
{
    local f="$1"

    if _is_symlink "$f"
    then
        echo "Error: Whitelist file is a symlink ($f)" >&2
        exit 1
    fi

    if [ ! -f "$f" ]
    then
        echo "Error: Whitelist file is not a regular file ($f)" >&2
        exit 1
    fi

    local uid mode
    uid="$(stat -c %u "$f" 2>/dev/null)" || { echo "Error: Cannot stat $f" >&2; exit 1; }
    mode="$(stat -c %a "$f" 2>/dev/null)" || { echo "Error: Cannot stat $f" >&2; exit 1; }

    if [ "$uid" != "0" ]
    then
        echo "Error: Whitelist file owner is not root ($f, uid=$uid)" >&2
        exit 1
    fi

    # Reject group/world writable files
    case "$mode" in
        ???)
            local g="${mode#?}"; g="${g%?}"
            local o="${mode#??}"
            case "$g" in 2|3|6|7) echo "Error: Whitelist file is group-writable ($f, mode=$mode)" >&2; exit 1 ;; esac
            case "$o" in 2|3|6|7) echo "Error: Whitelist file is world-writable ($f, mode=$mode)" >&2; exit 1 ;; esac
        ;;
        *)
            echo "Error: Whitelist file mode unreadable ($f, mode=$mode)" >&2
            exit 1
        ;;
    esac
}

function load_whitelist_file()
{
    local f="$1"
    whitelist_file_check "$f"

    local ip

    while IFS= read -r ip
    do
        [ -n "$ip" ] || continue

        if ! _is_ip_or_cidr "$ip"
        then
            echo "Error: invalid IP/CIDR in whitelist file $f: $ip" >&2
            exit 1
        fi

        case "$ip" in
            *:*)
                ipset add "$WHITE_V6" "$ip" -exist 2>/dev/null || { echo "Error: failed to add to $WHITE_V6: $ip" >&2; exit 1; }
            ;;
            *)
                ipset add "$WHITE_V4" "$ip" -exist 2>/dev/null || { echo "Error: failed to add to $WHITE_V4: $ip" >&2; exit 1; }
            ;;
        esac
    done < <(awk '
        /^[[:space:]]*$/ { next }
        /^[[:space:]]*#/ { next }
        {
            sub(/[[:space:]]*#.*/, "", $0)
            gsub(/[[:space:]]+/, "", $0)
            if ($0 != "") print $0
        }
    ' "$f")
}

function load_whitelist_dir()
{
    local d="$1"
    [ -n "$d" ] || return 0
    [ -d "$d" ] || return 0

    whitelist_dir_check "$d"

    shopt -s nullglob
    local f
    for f in "$d"/$WHITELIST_GLOB
    do
        load_whitelist_file "$f"
    done
    shopt -u nullglob
}

function list_active_blocks()
{
    echo "--- SxWall Active Configuration ---"

    local man_v4=0
    local man_v6=0
    local pre_v4=0
    local pre_v6=0

    if ipset list "$MANUAL_V4" >/dev/null 2>&1
    then
        man_v4="$(ipset list "$MANUAL_V4" 2>/dev/null | awk '/^Number of entries:/ {print $4; exit} END { }')"
        man_v4="${man_v4:-0}"
    fi

    if ipset list "$MANUAL_V6" >/dev/null 2>&1
    then
        man_v6="$(ipset list "$MANUAL_V6" 2>/dev/null | awk '/^Number of entries:/ {print $4; exit} END { }')"
        man_v6="${man_v6:-0}"
    fi

    if [ -n "${PREDEF_V4:-}" ] && ipset list "$PREDEF_V4" >/dev/null 2>&1
    then
        pre_v4="$(ipset list "$PREDEF_V4" 2>/dev/null | awk '/^Number of entries:/ {print $4; exit} END { }')"
        pre_v4="${pre_v4:-0}"
    fi

    if [ -n "${PREDEF_V6:-}" ] && ipset list "$PREDEF_V6" >/dev/null 2>&1
    then
        pre_v6="$(ipset list "$PREDEF_V6" 2>/dev/null | awk '/^Number of entries:/ {print $4; exit} END { }')"
        pre_v6="${pre_v6:-0}"
    fi

    echo "Manual Blocks:    IPv4: $man_v4 entries | IPv6: $man_v6 entries"
    
    if [ -n "${PREDEF_V4:-}" ] || [ -n "${PREDEF_V6:-}" ]
    then
        echo "Predef Blacklist: IPv4: $pre_v4 entries | IPv6: $pre_v6 entries"
    fi
    
    echo

    local members countries asns
    members="$(ipset list "$MASTER_V4" 2>/dev/null | awk '
        $1=="Members:" {inm=1; next}
        inm==1 && NF==1 {print $1}
    ')"

    countries="$(printf '%s\n' "$members" | awk '
        /^sx_c_[a-z0-9]+_v4$/ {
            s=$0
            sub(/^sx_c_/, "", s)
            sub(/_v4$/, "", s)
            print s
        }
    ' | sort -u)"

    asns="$(printf '%s\n' "$members" | awk '
        /^sx_a_[0-9]+_v4$/ {
            s=$0
            sub(/^sx_a_/, "", s)
            sub(/_v4$/, "", s)
            print "AS" s
        }
    ' | sort -u)"

    local zones_count=0
    local countries_line="None"

    if [ -n "$countries" ]
    then
        zones_count="$(printf '%s\n' "$countries" | wc -l | tr -d ' ')"
        countries_line="$(printf '%s\n' "$countries" | tr '\n' ' ' | sed 's/[[:space:]]*$//')"
    fi
    
    echo "Active Countries ($zones_count): $countries_line"

    local asn_count=0
    local asns_line="None"
    
    if [ -n "$asns" ]
    then
        asn_count="$(printf '%s\n' "$asns" | wc -l | tr -d ' ')"
        asns_line="$(printf '%s\n' "$asns" | tr '\n' ' ' | sed 's/[[:space:]]*$//')"
    fi
    
    echo "Active ASNs ($asn_count):      $asns_line"

    echo "-----------------------------------"
}


function lookup_ip_asn() 
{
    local ip="$1"

    if ! _is_ip_or_cidr "$ip" || [ "${ip#*/}" != "$ip" ]
    then
    		echo "Error: Invalid IP (note: no CIDR)."
    		exit 1
    fi
    
    echo "Querying Team Cymru for $ip..."
    local result=$(whois -h whois.cymru.com " -v $ip" | tail -n 1)
    
    if [[ -z "$result" ]] 
    then 
	    echo "Error: No info found." 
	    exit 1 
    fi
    
    local asn=$(echo "$result" | awk -F'|' '{print $1}' | tr -d '[:space:]')
    local range=$(echo "$result" | awk -F'|' '{print $3}' | tr -d '[:space:]')
    local cc=$(echo "$result" | awk -F'|' '{print $4}' | tr -d '[:space:]')
    local owner=$(echo "$result" | awk -F'|' '{print $7}' | sed 's/^ *//')
    
    echo "------------------------------------------------"
    echo "Target:  $ip" ; echo "ASN:     AS$asn" ; echo "Range:   $range" ; echo "Country: $cc" ; echo "Owner:   $owner"
    echo "------------------------------------------------"
    echo "Command to block:  ${0##*/} block AS$asn"
}

function sanitize_runtime() 
{
    local setname="$1"
    
    iptables -F SXWALL 2>/dev/null
    while iptables -D INPUT -m set --match-set "$setname" src -j DROP 2>/dev/null; do :; done
    while iptables -D INPUT -m set --match-set "$setname" src -j ACCEPT 2>/dev/null; do :; done
    
    ip6tables -F SXWALL 2>/dev/null
    while ip6tables -D INPUT -m set --match-set "$setname" src -j DROP 2>/dev/null; do :; done
    while ip6tables -D INPUT -m set --match-set "$setname" src -j ACCEPT 2>/dev/null; do :; done
}

function ensure_set_type() 
{
    local name="$1" ; local type="$2" ; local args="$3"
    if ipset -L "$name" >/dev/null 2>&1 
    then
        local cur_type=$(ipset -L "$name" | grep "Type:" | awk '{print $2}')
        if [[ "$cur_type" != "$type" ]] 
	then
            echo "Auto-Correction: Set $name is '$cur_type', expected '$type'. Correcting..."
            sanitize_runtime "$name"
            ipset destroy "$name"
            ipset create "$name" "$type" $args -exist
            echo "  -> Fixed."
        fi
    else
        ipset create "$name" "$type" $args -exist
    fi
}

function check_environment() 
{
    
    LOGDIR="/var/log/sxwall"
    mkdir -p "$LOGDIR" 2>/dev/null || true
    chmod 0750 "$LOGDIR" 2>/dev/null || true
    if [ -z "$SXWALL_LOGGING" ]
    then
    	LOGFILE="$LOGDIR/sxwall-$(date +"%Y-%m-%d-%H-%M-%S").log"
	SXWALL_LOGGING=1
    	exec > >(tee -a "$LOGFILE") 2>&1
    fi

    [[ ! -d "$DATADIR" ]] && mkdir -p "$DATADIR"
    [[ ! -d "$BACKUPDIR" ]] && mkdir -p "$BACKUPDIR"

    ensure_set_type "$WHITE_V4" "hash:net" "family inet"
    ensure_set_type "$WHITE_V6" "hash:net" "family inet6"
    
    ensure_set_type "$MASTER_V4" "list:set" "size 65535"
    ensure_set_type "$MASTER_V6" "list:set" "size 65535"
    
    ensure_set_type "$MANUAL_V4" "hash:net" "maxelem 1000000 family inet"
    ensure_set_type "$MANUAL_V6" "hash:net" "maxelem 1000000 family inet6"
    
    ensure_set_type "$PREDEF_V4" "hash:net" "maxelem 1000000 family inet"
    ensure_set_type "$PREDEF_V6" "hash:net" "maxelem 1000000 family inet6"

    ipset add "$MASTER_V4" "$MANUAL_V4" -exist
    ipset add "$MASTER_V6" "$MANUAL_V6" -exist

    ipset add "$MASTER_V4" "$PREDEF_V4" -exist
    ipset add "$MASTER_V6" "$PREDEF_V6" -exist

    ipset flush "$WHITE_V4" ; ipset flush "$WHITE_V6"
    ipset flush "$PREDEF_V4" ; ipset flush "$PREDEF_V6"

    for ip in "${WHITELIST[@]}"
    do
	ip="${ip//[[:space:]]/}"
	[ -z "$ip" ] && continue
    	
	if ! _is_ip_or_cidr "$ip"
    	then
        	echo "Error: invalid IP/CIDR in WHITELIST: $ip" >&2
        	exit 1
    	fi

    	case "$ip" in
        	*:*)
        		ipset add "$WHITE_V6" "$ip" -exist 2>/dev/null || { echo "Error: failed to add to $WHITE_V6: $ip" >&2; exit 1; }
			;;
        	
		*)
        		ipset add "$WHITE_V4" "$ip" -exist 2>/dev/null || { echo "Error: failed to add to $WHITE_V4: $ip" >&2; exit 1; }
			;;
    	esac
    done

    # External blacklist directory

    if [ -n "$BLACKLIST_DIR" ]
    then
        if [ ! -e "$BLACKLIST_DIR" ]
        then
            mkdir -p "$BLACKLIST_DIR" 2>/dev/null || { echo "Error: Cannot create $BLACKLIST_DIR" >&2; exit 1; }
            chown root:root "$BLACKLIST_DIR" 2>/dev/null || true
            chmod 0750 "$BLACKLIST_DIR" 2>/dev/null || true
        fi

        if [ -d "$BLACKLIST_DIR" ]
        then
            load_blacklist_dir "$BLACKLIST_DIR"
        else
            echo "Error: BLACKLIST_DIR exists but is not a directory ($BLACKLIST_DIR)" >&2
            exit 1
        fi
    fi

    # External whitelist directory 
    
        # Optional conf.d whitelist (auto-create)
    if [ -n "$WHITELIST_DIR" ]
    then
        if [ ! -e "$WHITELIST_DIR" ]
        then
            mkdir -p "$WHITELIST_DIR" 2>/dev/null || { echo "Error: Cannot create $WHITELIST_DIR" >&2; exit 1; }
            chown root:root "$WHITELIST_DIR" 2>/dev/null || true
            chmod 0750 "$WHITELIST_DIR" 2>/dev/null || true
        fi

        if [ -d "$WHITELIST_DIR" ]
        then
            load_whitelist_dir "$WHITELIST_DIR"
        else
            echo "Error: WHITELIST_DIR exists but is not a directory ($WHITELIST_DIR)" >&2
            exit 1
        fi
    fi
 
}

function inject_rules() 
{
    iptables -L SXWALL >/dev/null 2>&1 || { echo "Error: chain SXWALL missing"; exit 1; }
    iptables -C INPUT -j SXWALL >/dev/null 2>&1 || { echo "Error: INPUT not hooked to SXWALL"; exit 1; }

    iptables -F SXWALL 2>/dev/null
    iptables -A SXWALL -m set --match-set "$WHITE_V4" src -j ACCEPT
    iptables -A SXWALL -m set --match-set "$MASTER_V4" src -j DROP
    
    if ip6tables -F SXWALL 2>/dev/null 
    then
        ip6tables -A SXWALL -m set --match-set "$WHITE_V6" src -j ACCEPT 2>/dev/null
        ip6tables -A SXWALL -m set --match-set "$MASTER_V6" src -j DROP 2>/dev/null
    fi
}

function clean_all_blocks()
{
    echo "!!! CLEANING ALL BLOCKS !!!"

    ipset flush "$MASTER_V4" -exist
    ipset flush "$MASTER_V6" -exist

    ipset flush "$MANUAL_V4" -exist
    ipset flush "$MANUAL_V6" -exist

    ipset flush "$PREDEF_V4" -exist 2>/dev/null || true
    ipset flush "$PREDEF_V6" -exist 2>/dev/null || true

    ipset add "$MASTER_V4" "$MANUAL_V4" -exist
    ipset add "$MASTER_V6" "$MANUAL_V6" -exist
    ipset add "$MASTER_V4" "$PREDEF_V4" -exist
    ipset add "$MASTER_V6" "$PREDEF_V6" -exist

    local sets
    sets="$(ipset list -n 2>/dev/null | grep -E '^(sx_c_|sx_a_)')"
    for s in $sets
    do
        ipset destroy "$s" 2>/dev/null
    done

    check_environment

    echo "All blocks removed. Whitelist preserved. Firewall is now open."
}

function backup_sets() 
{
    local ts=$(date +"%Y-%m-%d-%H-%M") ; local fp="$BACKUPDIR/sxwall.conf-$ts"
    ipset save > "$fp" ; ln -sf "$fp" "$BACKUPDIR/latest.conf"
    [[ -s "$fp" ]] && echo "Saved to $fp" || rm -f "$fp"
}

function restore_sets() 
{
    local tgt="$1" ; local fp=""
    [[ "$tgt" == "latest" ]] && fp="$BACKUPDIR/latest.conf"
    [[ -f "$tgt" ]] && fp="$tgt"
    [[ -f "$BACKUPDIR/$tgt" ]] && fp="$BACKUPDIR/$tgt"
    [[ -f "$BACKUPDIR/sxwall.conf-$tgt" ]] && fp="$BACKUPDIR/sxwall.conf-$tgt"
    
    if [[ -f "$fp" ]] 
    then 
	    echo "Restoring sets..."  
	    ipset restore -! < "$fp"
    else 
	    echo "Warning: No backup found."  
    fi
}

function fetch_asn_data() 
{
    local input_asn="$1" ; local asn=""
    
    if [[ "$input_asn" =~ ^[0-9]+$ ]] 
    then 
	    asn="AS${input_asn}" 
    else 
	    asn="${input_asn}" 
    fi
    
    local raw="$DATADIR/${asn}.raw" ; local st="$DATADIR/${asn}.new" ; local fn="$DATADIR/${asn}.zone"
    
    echo -ne "Fetching $asn... \r"
    whois -h whois.radb.net -- "-i origin $asn" | grep ^route | sed 's/[ \t]//g' | sed 's/route6://g' | sed 's/route://g' > "$raw"
    
    if [[ ! -s "$raw" ]] 
    then 
	    rm -f "$raw"  
	    echo "Error: Failed to fetch data for $asn" 
	    return 1  
    fi
    
    if command -v iprange &> /dev/null 
    then 
	    grep -v ':' "$raw" | iprange --merge > "$st"  
	    grep ':' "$raw" >> "$st"
    else 
	    cat "$raw" > "$st"  
    fi
    
    rm -f "$raw"
    
    [[ -s "$st" ]] && mv "$st" "$fn" && echo "Updated $asn" || rm -f "$st"
}

function update_countries() {
    echo '--- Updating Countries ---'
    
    for z in $ZONES_TO_UPDATE 
    do
        echo -ne "Zone $z... \r"
        local t4="$DATADIR/${z}_v4.tmp" ; local t6="$DATADIR/${z}_v6.tmp" ; local st="$DATADIR/${z}.new" ; local fn="$DATADIR/${z}.zone"
        wget -nv -T 15 -t 3 "$URL_V4/${z}-aggregated.zone" -O "$t4" 2>/dev/null
        wget -nv -T 15 -t 3 "$URL_V6/${z}-aggregated.zone" -O "$t6" 2>/dev/null
        
	if [[ -s "$t4" ]] || [[ -s "$t6" ]] 
	then 
		cat "$t4" 2>/dev/null > "$st" 
		cat "$t6" 2>/dev/null >> "$st" 
		[[ -s "$st" ]] && mv "$st" "$fn" 
	fi
        
	rm -f "$t4" "$t6" "$st"
    done
    
    echo ''
}

function update_cached_asns() 
{ 
	shopt -s nullglob 
	
	for f in "$DATADIR"/AS*.zone
	do 
		fetch_asn_data "$(basename "$f" .zone)" 
	done  
	shopt -u nullglob 
}

function manage_manual_ip() 
{
    local act="$1" ; local ip="$2" 
    local sn="" ; local cmd=""
    
    [[ "$act" == "block" ]] && cmd="add" || cmd="del"
    
    if [[ "$ip" =~ : ]] 
    then 
	    sn="$MANUAL_V6" 
    else 
	    sn="$MANUAL_V4" 
    fi
    
    ipset $cmd "$sn" "$ip" -exist 
    
    echo "Manual $act on $ip ($sn)"
}

function manage_subset() 
{
    local act="$1" ; local type="$2" ; local id="$3" ; local src="$4" ; local pfx="" 
    
    [[ "$type" == "country" ]] && pfx="sx_c_${id}" 
    [[ "$type" == "asn" ]] && pfx="sx_a_${id//AS/}"
    
    local sv4="${pfx}_v4" ; local sv6="${pfx}_v6" ; local tv4="${sv4}_tmp" ; local tv6="${sv6}_tmp"
    
    if [[ "$act" == "block" ]] 
    then
        echo "Blocking $id..."
        ipset create "$sv4" hash:net maxelem 2000000 family inet -exist
        ipset create "$tv4" hash:net maxelem 2000000 family inet -exist ; ipset flush "$tv4"
        grep -v ':' "$src" | sed "s/^/add $tv4 /" | ipset restore -!
        ipset swap "$tv4" "$sv4" ; ipset destroy "$tv4" ; ipset add $MASTER_V4 "$sv4" -exist 2>/dev/null

        ipset create "$sv6" hash:net maxelem 2000000 family inet6 -exist
        ipset create "$tv6" hash:net maxelem 2000000 family inet6 -exist ; ipset flush "$tv6"
        grep ':' "$src" | sed "s/^/add $tv6 /" | ipset restore -!
        ipset swap "$tv6" "$sv6" ; ipset destroy "$tv6" ; ipset add $MASTER_V6 "$sv6" -exist 2>/dev/null
    elif [[ "$act" == "unblock" ]] 
    then
        echo "Unblocking $id..."
        ipset del $MASTER_V4 "$sv4" -exist 2>/dev/null ; ipset destroy "$sv4" 2>/dev/null
        ipset del $MASTER_V6 "$sv6" -exist 2>/dev/null ; ipset destroy "$sv6" 2>/dev/null
    fi
}

function core() 
{
    local act="$1" ; local tgt="$2"

    if [[ "$tgt" =~ ^(AF|AS|EU|NA|SA|OC|AN)$ ]]
    then
        local cc
        local list
        list="$(region_to_countries "$tgt")"

        if [[ -z "$list" ]]
        then
            echo "Error: Unknown region."
            exit 1
        fi

        echo "Applying $act for region $tgt..."
        for cc in $list
        do
            if [[ ! -f "$DATADIR/$cc.zone" ]]
            then
                echo "Error: Zone $cc not found. Run update."
                exit 1
            fi
            manage_subset "$act" "country" "$cc" "$DATADIR/$cc.zone"
        done
        return 0
    fi
    
    if [[ "$tgt" =~ ^[a-z]{2}$ ]] 
    then 
        if [[ ! -f "$DATADIR/$tgt.zone" ]] 
	then 
		echo "Error: Zone $tgt not found. Run update." 
		exit 1 
	fi
        
	manage_subset "$act" "country" "$tgt" "$DATADIR/$tgt.zone"
    
    elif [[ "$tgt" =~ ^AS[0-9]+$ || "$tgt" =~ ^[0-9]+$ ]] 
    then 
        local asn_clean="" 
	if [[ "$tgt" =~ ^[0-9]+$ ]] 
	then 
		asn_clean="AS$tgt" 
	else 
		asn_clean="$tgt" 
	fi
        
	if [[ "$act" == "block" && ! -f "$DATADIR/$asn_clean.zone" ]] 
	then 
            fetch_asn_data "$tgt"
            if [[ ! -f "$DATADIR/$asn_clean.zone" ]] 
	    then 
		    echo 'Aborting: ASN data missing.'
		    exit 1 
	    fi
        fi
        
	manage_subset "$act" "asn" "$asn_clean" "$DATADIR/$asn_clean.zone"
    else 
	    manage_manual_ip "$act" "$tgt"  
    fi
}

function _ipset_exists()
{
    ipset list "$1" >/dev/null 2>&1
}

function _ipset_has_set_member()
{
    # $1 = master set (list:set), $2 = child set name
    ipset list "$1" 2>/dev/null | grep -qE "^[[:space:]]+$2$"
}

function _master_members()
{
    # Prints child set names referenced by a list:set
    ipset list "$1" 2>/dev/null | awk '
        $1=="Members:" {inm=1; next}
        inm==1 && NF==1 {print $1}
    '
}

function _check_ip()
{
    local ip="$1"
    local fam="v4"
    local wset="$WHITE_V4"
    local mset="$MASTER_V4"

    if _is_ipv6 "$ip"
    then
        fam="v6"
        wset="$WHITE_V6"
        mset="$MASTER_V6"
    fi

    echo "---- SxWall Check ($fam) ----"
    echo "Target: $ip"
    echo

    if ipset test "$wset" "$ip" >/dev/null 2>&1
    then
        echo "Status: ALLOWED (whitelisted: $wset)"
        return 0
    fi

    local man="$MANUAL_V4"
    local pre="$PREDEF_V4"
    if [[ "$fam" == "v6" ]]
    then
        man="$MANUAL_V6"
        pre="$PREDEF_V6"
    fi

    if ipset test "$man" "$ip" >/dev/null 2>&1
    then
        echo "Status: BLOCKED (manual: $man)"
        return 0
    fi

    if ipset test "$pre" "$ip" >/dev/null 2>&1
    then
        echo "Status: BLOCKED (predef: $pre)"
        return 0
    fi

    local hit=0
    local s
    while read -r s
    do
        [[ -z "$s" ]] && continue
        [[ "$s" == "$man" ]] && continue
        [[ "$s" == "$pre" ]] && continue

        if ipset test "$s" "$ip" >/dev/null 2>&1
        then
            if [[ $hit -eq 0 ]]
            then
                echo "Status: BLOCKED"
                echo "Matched by:"
            fi
            echo "  - $s"
            hit=1
        fi
    done < <(_master_members "$mset")

    if [[ $hit -eq 0 ]]
    then
        echo "Status: NOT BLOCKED"
    fi

    return 0
}

function _check_country()
{
    local cc="$1"
    local sv4="sx_c_${cc}_v4"
    local sv6="sx_c_${cc}_v6"

    echo "---- SxWall Check (country) ----"
    echo "Target: $cc"
    echo

    local active4="no"
    local active6="no"

    if _ipset_exists "$sv4" && _ipset_has_set_member "$MASTER_V4" "$sv4"
    then active4="yes"
    fi

    if _ipset_exists "$sv6" && _ipset_has_set_member "$MASTER_V6" "$sv6"
    then active6="yes"
    fi

    echo "Active: IPv4=$active4 | IPv6=$active6"

    if _ipset_exists "$sv4"
    then
        local n4
        n4="$(ipset list "$sv4" 2>/dev/null | awk '/Number of entries:/ {print $4}')"
        echo "Set: $sv4 entries=${n4:-unknown}"
    else
        echo "Set: $sv4 missing"
    fi

    if _ipset_exists "$sv6"
    then
        local n6
        n6="$(ipset list "$sv6" 2>/dev/null | awk '/Number of entries:/ {print $4}')"
        echo "Set: $sv6 entries=${n6:-unknown}"
    else
        echo "Set: $sv6 missing"
    fi

    local zonefile="$DATADIR/${cc}.zone"
    if [[ -f "$zonefile" ]]
    then
        echo "Zonefile: $zonefile (present)"
    else
        echo "Zonefile: $zonefile (missing)"
    fi
}

function _check_asn()
{
    local asn="$1"
    [[ "$asn" =~ ^[0-9]+$ ]] && asn="AS$asn"
    local clean="${asn//AS/}"

    local sv4="sx_a_${clean}_v4"
    local sv6="sx_a_${clean}_v6"

    echo "---- SxWall Check (asn) ----"
    echo "Target: $asn"
    echo

    local active4="no"
    local active6="no"

    if _ipset_exists "$sv4" && _ipset_has_set_member "$MASTER_V4" "$sv4"
    then active4="yes"
    fi

    if _ipset_exists "$sv6" && _ipset_has_set_member "$MASTER_V6" "$sv6"
    then active6="yes"
    fi

    echo "Active: IPv4=$active4 | IPv6=$active6"

    if _ipset_exists "$sv4"
    then
        local n4
        n4="$(ipset list "$sv4" 2>/dev/null | awk '/Number of entries:/ {print $4}')"
        echo "Set: $sv4 entries=${n4:-unknown}"
    else
        echo "Set: $sv4 missing"
    fi

    if _ipset_exists "$sv6"
    then
        local n6
        n6="$(ipset list "$sv6" 2>/dev/null | awk '/Number of entries:/ {print $4}')"
        echo "Set: $sv6 entries=${n6:-unknown}"
    else
        echo "Set: $sv6 missing"
    fi

    local zonefile="$DATADIR/${asn}.zone"
    if [[ -f "$zonefile" ]]
    then
        echo "Zonefile: $zonefile (present)"
    else
        echo "Zonefile: $zonefile (missing)"
    fi
}

function _check_region()
{
    local r="$1"
    local list=""
    case "$r" in
        AF) list="$AF" ;;
        AS) list="$AS" ;;
        EU) list="$EU" ;;
        NA) list="$NA" ;;
        SA) list="$SA" ;;
        OC) list="$OC" ;;
        AN) list="$AN" ;;
        *) echo "Error: Unknown region. Use AF|AS|EU|NA|SA|OC|AN"; exit 1 ;;
    esac

    echo "---- SxWall Check (region) ----"
    echo "Target: $r"
    echo

    local cc
    local active=""
    for cc in $list
    do
        local sv4="sx_c_${cc}_v4"
        local sv6="sx_c_${cc}_v6"
        if (_ipset_exists "$sv4" && _ipset_has_set_member "$MASTER_V4" "$sv4") || (_ipset_exists "$sv6" && _ipset_has_set_member "$MASTER_V6" "$sv6")
        then
            active="$active $cc"
        fi
    done

    if [[ -z "$active" ]]
    then
        echo "Active countries: none"
    else
        echo "Active countries:$active"
    fi
}

function check_target()
{
    local tgt="$1"

    if _is_ip_or_cidr "$tgt"
    then
        _check_ip "$tgt"
        return 0
    fi

    # Region
    if [[ "$tgt" =~ ^(AF|AS|EU|NA|SA|OC|AN)$ ]]
    then
        _check_region "$tgt"
        return 0
    fi

    # Country
    if [[ "$tgt" =~ ^[a-z]{2}$ ]]
    then
        _check_country "$tgt"
        return 0
    fi

    # ASN
    if [[ "$tgt" =~ ^AS[0-9]+$ || "$tgt" =~ ^[0-9]+$ ]]
    then
        _check_asn "$tgt"
        return 0
    fi

    echo "Error: Unsupported target for check."
    echo "Supported: IP/CIDR, country (fr), ASN (AS123 or 123), region (AF|AS|EU|NA|SA|OC|AN)"
    exit 1
}

function import_manual_list()
{
    local src="$1"
    local mode="$2"

    local in
    local buf=""
    local ip

    if [[ "$src" == "-" ]]
    then
        in="/dev/stdin"
    else
        in="$src"
        [[ -f "$in" ]] || { echo "Error: file not found: $in" >&2; exit 1; }
    fi

    if [[ "$mode" == "--clear" ]]
    then
        ipset flush "$MANUAL_V4" 2>/dev/null || true
        ipset flush "$MANUAL_V6" 2>/dev/null || true
    fi

    while IFS= read -r ip
    do
        ip="${ip%%#*}"
        ip="${ip//[[:space:]]/}"
        [ -z "$ip" ] && continue

        if ! _is_ip_or_cidr "$ip"
        then
            echo "Error: invalid IP/CIDR in import: $ip" >&2
            exit 1
        fi

        case "$ip" in
            *:*) buf+="add $MANUAL_V6 $ip -exist"$'\n' ;;
            *)   buf+="add $MANUAL_V4 $ip -exist"$'\n' ;;
        esac
    done < "$in"

    if [ -n "$buf" ]
    then
        printf '%s' "$buf" | ipset restore || { echo "Error: import failed (source: $src)." >&2; exit 1; }
    fi

    echo "Import done."
}

[[ -z "$1" ]] && show_usage

case "$ACTION" in
    start)
        echo "--- SxWall Startup ---"
        
	if [[ -f "$IPTABLES_FILE" ]] 
	then 
		iptables-restore < "$IPTABLES_FILE" 
	fi
        
	restore_sets "latest" ; check_environment ; inject_rules
        
	echo "$SCRIPT_NAME start: Initialize firewall (loads WHITELIST/BLACKLIST + $WHITELIST_DIR/$WHITELIST_GLOB - $BLACKLIST_DIR/$BLACKLIST_GLOB)"
        echo "--- Ready ---" 
    	;;
  
    reload)
        check_environment ; inject_rules
        ;;

    update)
	if [[ "$TARGET" == "asn" && -n "$PARAM" ]] 
	then 
		fetch_asn_data "$PARAM"
        else 
		update_countries 
		update_cached_asns 
	fi 
    	;;
    
    save) 
	backup_sets 	    
    	;;

   restore)
        [[ -z "$TARGET" ]] && show_usage
        restore_sets "$TARGET" ; check_environment ; inject_rules
        ;;
 
    whoas) 
	[[ -z "$TARGET" ]] && echo "Error: Missing IP" && exit 1 
	lookup_ip_asn "$TARGET" 
	;;
    
    import)
        [[ -z "$TARGET" ]] && show_usage
        check_environment
        import_manual_list "$TARGET" "$PARAM"
        inject_rules
        ;;
    
    list) 
	list_active_blocks 
	;;
    
    clean) 
	clean_all_blocks 
	;;
    
    check)
        [[ -z "$TARGET" ]] && show_usage
        check_environment
        check_target "$TARGET"
        ;;
 
    block|unblock) 
        [[ -z "$TARGET" ]] && show_usage
        check_environment 
	core "$ACTION" "$TARGET" 
	inject_rules 
	;;
    
    help|*) 
	    show_usage 
	    ;;
esac