#!/bin/bash

# Get all open ports on the server

json_file="/usr/local/x-ui/bin/config.json"

IFS=$'\n'

ports=($(jq -r '.inbounds[].port' "$json_file"))

#echo $ports

while true; do
	rules=$(sudo iptables -L INPUT -n --line-numbers)
	connections=$(sudo netstat -tn)
	for port in "${ports[@]}"; do
		if [[ ! -z "$port" ]] && [[ "$port" -ge 10000 ]]; then
			# Get connected IPs on the specific port
			connectedRules=$(echo "$rules" | grep ":$port")
			connectedIPs=$(echo "$connections" | grep ":$port" | awk '{print $5}' | cut -d: -f1 | sort | uniq)
			whitelistedIP=""
			ipCount=0

			IFS=$'\n'
			for ip in $(echo "$connectedIPs"); do
				ip=$(echo "$ip" | tr -d '[:space:]')
				if [[ ! -z "$ip" ]]; then
					((ipCount++))
					if [[ "$ipCount" -eq 1 ]]; then
						whitelistedIP=$ip
					fi
				fi
			done
			#if any ip has even been whited on this port
			whited=$(echo "$connectedRules" | grep "$whitelistedIP.*:$port" | sort -r)
			# if there was reject rule applied for this port or not
			has_reject=$(echo "$connectedRules" | grep "REJECT.*:$port" | sort -r)
			ruleNumbers=$(echo "$connectedRules" | awk '{print $1}' | sort -r)
			#if any ip connected run the process
			if [[ "$ipCount" -gt 0 ]]; then
				#echo "whitelist ip is $whitelistedIP"
				#echo $connectedRules
				#echo $whited
				#echo $has_reject
				#echo $ruleNumbers
				if [[ -z "$has_reject" ]]; then
					#echo "has reject!"
					sudo iptables -A INPUT -p tcp --dport "$port" -j REJECT --reject-with tcp-reset
				fi
				if [[ -z "$whited" ]]; then
					#echo "has whited!"
					sudo iptables -I INPUT -p tcp --dport "$port" -s "$whitelistedIP" -j ACCEPT
				fi
			#otherwise clean everything
			else
				#echo "all ips cleared!"
				# Clear restrictions if no IP is connected
				for ruleNumber in $ruleNumbers; do
					sudo iptables -D INPUT "$ruleNumber"
				done
			fi
		fi
	done
	sleep 5
done