#!/bin/sh

#### Variables ####
iproute=$(ip route get 2001:4860:4860::8888 | awk -- '{print $7}')
seedlist_dir=/var/lib/ipset
blists=$(find "$seedlist_dir" -name "*-inet6.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 ]
        ip6tables -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 f in $blists
    do
        listdir=$(basename -s ".save" "$f")
        #### Check for SeedList Dir ####
        if [ ! -d "$seedlist_dir" ]; then
            echo "ERROR: Seedlist directory does not exist" >&2
            exit 2
        fi
        
        ipset restore -! < "$f"
        
        for i in $listdir
            do
                if chain_exists "$i"-input; then
                    ip6tables -D INPUT -i "$iproute" -m set --match-set "$listdir" src -j "$listdir"-input || true
                    ip6tables -D "$listdir"-input -i "$iproute" -j DROP -m comment --comment "$listdir"-input || true
                    ip6tables -F "$listdir"-input
                    ip6tables -X "$listdir"-input
                fi
            
                ip6tables -N "$listdir"-input
                ip6tables -A "$listdir"-input -i "$iproute" -j DROP -m comment --comment "$listdir"-input
                ip6tables -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
                    ip6tables -D OUTPUT -o "$iproute" -m set --match-set "$listdir" dst -j "$listdir"-output || true
                    ip6tables -D "$listdir"-output -j DROP -m comment --comment "$listdir"-output || true
                    ip6tables -F "$listdir"-output
                    ip6tables -X "$listdir"-output
                fi
                
                ip6tables -N "$listdir"-output
                ip6tables -A "$listdir"-output -j DROP -m comment --comment "$listdir"-output
                ip6tables -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
                    ip6tables -D FORWARD -o "$iproute" -m set --match-set "$listdir" dst -j "$listdir"-forward || true
                    ip6tables -D "$listdir"-forward -o "$iproute" -j DROP -m comment --comment "$listdir"-forward || true
                    ip6tables -F "$listdir"-forward
                    ip6tables -X "$listdir"-forward
                fi
                
                ip6tables -N "$listdir"-forward
                ip6tables -A "$listdir"-forward -o "$iproute" -j DROP -m comment --comment "$listdir"-forward
                ip6tables -I FORWARD -o "$iproute" -m set --match-set "$listdir" dst -j "$listdir"-forward
            done
    done
    ;;
stop)
    for f in $blists
        do
            listdir=$(basename -s ".save" "$f")
            #### 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
                    ip6tables -D INPUT -i "$iproute" -m set --match-set "$listdir" src -j "$listdir"-input || true
                    ip6tables -D "$listdir"-input -i "$iproute" -j DROP -m comment --comment "$listdir"-input || true
                    ip6tables -F "$listdir"-input
                    ip6tables -X "$listdir"-input
                fi
            done

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

            for i in $listdir
            do
                if chain_exists "$i"-forward; then
                    ip6tables -D FORWARD -o "$iproute" -m set --match-set "$listdir" dst -j "$listdir"-forward || true
                    ip6tables -D "$listdir"-forward -o "$iproute" -j DROP -m comment --comment "$listdir"-forward || true
                    ip6tables -F "$listdir"-forward
                    ip6tables -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 ip6tables block/byte counts ####
    echo ''
    echo 'Total Block/Byte Counts by List Inbound'
    echo ''
    ip6tables -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 ''
    ip6tables -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 ''
    ip6tables -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 ip6tables blocklist counts ####
    ipi=$( ip6tables -L INPUT -nvx --line-numbers | grep bl-| awk '{print $1}')
    for i in $ipi
        do 
            ip6tables -Z INPUT "$i"
        done

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