#!/bin/bash

##########################################################
###    bash <(curl -fsSL https://raw.lhy.life/nft.sh) --port_allow=80,443,11000-12000
##########################################################

## GetArgValue <key> <defValue>
## ./script --key=value
g_args=("$@")
function GetArgValue()
{
	local key="$1"
	local defvalue="$2"

	for arg in "${g_args[@]}"
	do
		if [[ "$arg" == "--$key="* ]]
		then
			echo "${arg#*=}"
			return
		fi
	done

	echo "$defvalue"
}


PORT_ALLOW=$(GetArgValue "port_allow" "80,443,22")

## Censys
IP_REJECT='66.132.159.0/24, 162.142.125.0/24, 167.94.138.0/24, 167.94.145.0/24, 167.94.146.0/24, 167.248.133.0/24, 199.45.154.0/24, 199.45.155.0/24, 206.168.34.0/24, 206.168.35.0/24' 
IP6_REJECT='2602:80d:1000:b0cc:e::/80, 2620:96:e000:b0cc:e::/80, 2602:80d:1003::/112, 2602:80d:1004::/112'

SELF_ALLOW=""
NEIGHBOR_REJECT=""
PUBLIC_IP=$(curl -4 -s ip.sb)

if [ -n "$PUBLIC_IP" ]; then
    SELF_ALLOW="ip saddr $PUBLIC_IP accept"
    NEIGHBOR_RANGE="${PUBLIC_IP%.*}.0/24"
    NEIGHBOR_REJECT="ip saddr $NEIGHBOR_RANGE log prefix \"Neighbors Reject: \" drop"
fi


cp /etc/nftables.conf /etc/nftables.conf.bak

cat <<EOF > /etc/nftables.conf
#!/usr/sbin/nft -f

flush ruleset

table inet main {
    chain input {
        type filter hook input priority filter; policy drop;
        iifname "lo" accept
        ct state established,related accept
        fib saddr type local accept
        $SELF_ALLOW
        $NEIGHBOR_REJECT
        ip saddr { $IP_REJECT } drop
        ip6 saddr { $IP6_REJECT } drop
        meta l4proto ipv6-icmp accept
        meta l4proto { tcp, udp } th dport { $PORT_ALLOW } accept
    }
}
EOF

systemctl restart nftables
systemctl enable nftables
nft list ruleset