#!/bin/bash

#### Variables ####
iproute=$(ip route get 8.8.8.8 | awk -- '{print $5}')
seedlist_dir=/var/lib/ipset
blists=$(find "$seedlist_dir" -name "*-inet.save")

#### Check for IPset ###
if [ ! "$(command -v /usr/sbin/ipset)" ]; then
    echo "ERROR: ipset binary not found in path"
    exit 2
fi

chain_exists() {
    {   
        [ $# -lt 1 ] || [ $# -gt 2 ] && { 
            echo "Usage: chain_exists <chain_name>" >&2
            exit 1
        }
        
        chain_name="$1"  
        shift
        [ $# -eq 1 ]
        iptables -n -L "$chain_name" >/dev/null 2>&1
    }
}

list_exists() {
    {   
        [ $# -ne 1 ] && { 
            echo "Usage: list_exists <list_name>" >&2
            exit 1
        }
        
        list_name="$1"
        ipset list "$list_name" -name >/dev/null 2>&1
    }
}

#### Command Line Options ####
case "$1" in
start)
for sl in $blists
    do
        listdir=$(basename -s ".save" "$sl")
        #### Check for SeedList Dir ####
        if [ ! -d "$seedlist_dir" ]; then
            echo "ERROR: Seedlist directory does not exist" >&2
            exit 2
        fi

        ipset restore -! < "$sl"

        for i in $listdir
            do
                if chain_exists "$i"-input; then
                    iptables -D INPUT -i "$iproute" -m set --match-set "$listdir" src -j "$listdir"-input || true
                    iptables -D "$listdir"-input -i "$iproute" -j DROP -m comment --comment "$listdir"-input || true
                    iptables -F "$listdir"-input
                    iptables -X "$listdir"-input
                fi

                iptables -N "$listdir"-input
                iptables -A "$listdir"-input -i "$iproute" -j DROP -m comment --comment "$listdir"-input
                iptables -I INPUT -i "$iproute" -m set --match-set "$listdir" src -j "$listdir"-input
            done

        for i in $listdir
            do
                if chain_exists "$i"-output; then
                    iptables -D OUTPUT -o "$iproute" -m set --match-set "$listdir" dst -j "$listdir"-output || true
                    iptables -D "$listdir"-output -j DROP -m comment --comment "$listdir"-output || true
                    iptables -F "$listdir"-output
                    iptables -X "$listdir"-output
                fi

                iptables -N "$listdir"-output
                iptables -A "$listdir"-output -j DROP -m comment --comment "$listdir"-output
                iptables -I OUTPUT -o "$iproute" -m set --match-set "$listdir" dst -j "$listdir"-output
            done        

        for i in $listdir
            do
                if chain_exists "$i"-forward; then
                    iptables -D FORWARD -o "$iproute" -m set --match-set "$listdir" dst -j "$listdir"-forward || true
                    iptables -D "$listdir"-forward -i "$iproute" -j DROP -m comment --comment "$listdir"-forward || true
                    iptables -F "$listdir"-forward
                    iptables -X "$listdir"-forward
                fi

                iptables -N "$listdir"-forward
                iptables -A "$listdir"-forward -i "$iproute" -j DROP -m comment --comment "$listdir"-forward
                iptables -I FORWARD -o "$iproute" -m set --match-set "$listdir" dst -j "$listdir"-forward
            done
    done
    ;;
stop)
    for sl in $blists
        do
            listdir=$(basename -s ".save" "$sl")
            #### Check for SeedList Dir ####
            if [ ! -d "$seedlist_dir" ]; then
                echo "ERROR: ipset data directory does not exist" >&2
                exit 2
            fi

            for i in $listdir
            do
                if chain_exists "$i"-input; then
                    iptables -D INPUT -i "$iproute" -m set --match-set "$listdir" src -j "$listdir"-input || true
                    iptables -D "$listdir"-input -i "$iproute" -j DROP -m comment --comment "$listdir"-input || true
                    iptables -F "$listdir"-input
                    iptables -X "$listdir"-input
                fi
            done

            for i in $listdir
            do
                if chain_exists "$i"-output; then
                    iptables -D OUTPUT -o "$iproute" -m set --match-set "$listdir" dst -j "$listdir"-output || true
                    iptables -D "$listdir"-output -o "$iproute"-j DROP -m comment --comment "$listdir"-output || true
                    iptables -F "$listdir"-output
                    iptables -X "$listdir"-output
                fi    
            done

            for i in $listdir
            do
                if chain_exists "$i"-forward; then
                    iptables -D FORWARD -o "$iproute" -m set --match-set "$listdir" dst -j "$listdir"-forward || true
                    iptables -D "$listdir"-forward -o "$iproute" -j DROP -m comment --comment "$listdir"-forward || true
                    iptables -F "$listdir"-forward
                    iptables -X "$listdir"-forward
                fi
            done

            for i in $listdir
            do
                if list_exists "$i"; then
                    ipset flush "$listdir"
                    ipset destroy -q "$listdir" || true
                fi
            done
        done
    ;;
status)
    echo '= after.init Blocklist(s) Status ='
    ipset -t list
    # show iptables block/byte counts
    echo ''
    echo 'Total Block/Byte Counts by List Inbound'
    echo ''
    iptables -L -nvx | grep "$listdir" | grep 'match-set' | grep input | awk 'BEGIN{print "Pkts Bytes Blocklist Prot Opt In Out Source Destination"}1' | column -t
    echo ""
    echo 'Total Block/Byte Counts by List Outbound'
    echo ''
    iptables -L -nvx | grep "$listdir" | grep 'match-set' | grep output | awk 'BEGIN{print "Pkts Bytes Blocklist Prot Opt In Out Source Destination"}1' | column -t
    echo ""
    echo 'Total Block/Byte Counts by List Forwarded'
    echo ''
    iptables -L -nvx | grep "$listdir" | grep 'match-set' | grep forward | awk 'BEGIN{print "Pkts Bytes Blocklist Prot Opt In Out Source Destination"}1' | column -t
    ;;
reset-count)
    echo '= after.init is resetting the Block/Byte Counts ='

    # reset iptables blocklist pkts/bytes
    ipi=$( iptables -L INPUT -nvx --line-numbers | grep bl-| awk '{print $1}')
    for i in $ipi
        do 
            iptables -Z INPUT "$i"
        done

    ipo=$( iptables -L OUTPUT -nvx --line-numbers | grep bl- | awk '{print $1}')
    for i in $ipo
        do
            iptables -Z OUTPUT "$i"
        done 
    
    ipf=$( iptables -L FORWARD -nvx --line-numbers | grep bl-| awk '{print $1}')
    for i in $ipf
        do
            iptables -Z FORWARD "$i"
        done
    ;;
*)
    echo "'$1' is not supported"
    echo 'Usage: /etc/ufw/after.init {start|stop|status|reset-count}'
    ;;
esac