#!/bin/bash

config_file="/etc/iplimit/iplimit.conf"

logger "iplimit service started!"
echo "iplimit service started!"

json_file="/usr/local/x-ui/bin/config.json"
IFS=$'\n'

if [ -f "$config_file" ]; then
	source "$config_file"
	source /etc/iplimit/db.sh
	
	ports=($(jq -r '.inbounds[].port' "$json_file"))

	iptables -F && iptables -t mangle -F
	echo "deleting root qdisc tc qdisc del dev "$interface" root"
	tc qdisc del dev "$interface" root
	tc qdisc del dev "$interface" handle ffff: ingress
	tc filter del dev "$interface"
	
	echo "done deleting"
	
	echo "$type"
	if [[ "$type" == "ip" ]]; then
		logger "running by IP limit!"
		echo "running by IP limit!"
		while true; do
			connections=$(netstat -tn | awk '{split($4, lparts, ":"); split($5, rparts, ":"); ip=rparts[1] ":" lparts[2]; if (!seen[ip]++) print ip;}')
			rules=$(sudo iptables -L INPUT -n --line-numbers)
			ports_setting=$(get_ports | grep -v '^$')
			for port in "${ports[@]}"; do
				connected_ips=$(echo "$connections" | grep ":$port" | grep -v '^$' | sort -n)
				connected_rules=$(echo "$rules" | grep ":$port")
				count=$(echo "$connected_ips" | wc -l)
				port_setting=$(echo "$ports_setting" | grep "$port")
				port_ignore=$(echo "$port_setting" | awk -F ',' '{print $4}' | tr -d '[:space:]')
				port_limit=$(echo "$port_setting" | awk -F ',' '{print $2}'| tr -d '[:space:]')
				if [ -n "$port_limit" ]; then
					max_conn="$port_limit"
				fi
				
				if [ "$count" -gt "$max_conn" ] && { [ -z "$port_ignore" ] || [ "$port_ignore" -ne 1 ]; }; then
				#if [ "$count" -gt "$max_conn" ]; then
					allows=$(echo "$connected_ips" | head -n "$max_conn" | awk -F: '{print $1}')
					if [[ -n "$allows" ]]; then
						for ip in $allows; do
							whited=$(echo "$connected_rules" | grep "$ip.*:$port" | sort -r)
							if [[ -z "$whited" ]]; then
								#echo "has whited!"
								sudo iptables -I INPUT -p tcp --dport "$port" -s "$ip" -j ACCEPT
							fi
						done
					fi
					#has_reject=$(echo "$connected_rules" | grep "DROP.*:$port" | sort -r)
					has_reject=$(echo "$connected_rules" | grep "DROP" | sort -r)
					if [[ -z "$has_reject" ]]; then
						#echo "has reject!"
						sudo iptables -A INPUT -p tcp --dport "$port" -j DROP #--reject-with tcp-reset
					fi
				else
					rules_number=$(echo "$connected_rules" | awk '{print $1}' | sort -r)
					for rule_number in $rules_number; do
						sudo iptables -D INPUT "$rule_number"
					done
				fi
			done
			sleep 10
		done;
	elif [[ "$type" == "port" ]]; then
		logger "running by PORT limit!"
		echo "running by PORT limit!"
		while true; do
			connections=$(netstat -tn | awk '{split($4, lparts, ":"); split($5, rparts, ":"); ip=rparts[1] ":" lparts[2]; if (!seen[ip]++) print ip;}')
			rules=$(sudo iptables -L INPUT -n --line-numbers)
			ports_setting=$(get_ports | grep -v '^$')
			for port in "${ports[@]}"; do
				connected_ips=$(echo "$connections" | grep ":$port" | grep -v '^$' | sort -n)
				connected_rules=$(echo "$rules" | grep ":$port")
				count=$(echo "$connected_ips" | wc -l)
				port_setting=$(echo "$ports_setting" | grep "$port")
				port_ignore=$(echo "$port_setting" | awk -F ',' '{print $4}' | tr -d '[:space:]')
				port_limit=$(echo "$port_setting" | awk -F ',' '{print $2}'| tr -d '[:space:]')
				if [ -n "$port_ignore" ] && [ "$port_ignore" -eq 1 ]; then
					continue
				fi
				if [ -n "$port_limit" ]; then
					max_conn="$port_limit"
				fi
				
				if [ "$count" -gt "$max_conn" ] && { [ -z "$port_ignore" ] || [ "$port_ignore" -ne 1 ]; }; then
				#if [ "$count" -gt "$max_conn" ]; then
					has_reject=$(echo "$connected_rules" | grep "DROP" | sort -r)
					#has_reject=$(echo "$connected_rules" | grep "DROP.*:$port" | sort -r)
					if [[ -z "$has_reject" ]]; then
						#has reject!
						sudo iptables -A INPUT -p tcp --dport "$port" -j DROP #--reject-with tcp-reset
					fi
				else
					rules_number=$(echo "$connected_rules" | awk '{print $1}' | sort -r)
					for rule_number in $rules_number; do
						sudo iptables -D INPUT "$rule_number"
					done
				fi
			done
			sleep 20
			iptables -F
			sleep 40
		done;
	elif [[ "$type" == "full" ]]; then
		logger "running by FULL limit!"
		echo "running by FULL limit!"
		while true; do
			connections=$(netstat -tn | awk '{split($4, lparts, ":"); split($5, rparts, ":"); ip=rparts[1] ":" lparts[2]; if (!seen[ip]++) print ip;}')
			rules=$(sudo iptables -L INPUT -n --line-numbers)
			ports_setting=$(get_ports | grep -v '^$')
			for port in "${ports[@]}"; do
				connected_ips=$(echo "$connections" | grep ":$port" | grep -v '^$' | sort -n)
				connected_rules=$(echo "$rules" | grep ":$port")
				count=$(echo "$connected_ips" | wc -l)
				port_setting=$(echo "$ports_setting" | grep "$port")
				port_ignore=$(echo "$port_setting" | awk -F ',' '{print $4}' | tr -d '[:space:]')
				port_limit=$(echo "$port_setting" | awk -F ',' '{print $2}'| tr -d '[:space:]')
				if [ -n "$port_limit" ]; then
					max_conn="$port_limit"
				fi
				
				if [ "$count" -gt "$max_conn" ] && { [ -z "$port_ignore" ] || [ "$port_ignore" -ne 1 ]; }; then
				#if [ "$count" -gt "$max_conn" ]; then
					allows=$(echo "$connected_ips" | head -n "$max_conn" | awk -F: '{print $1}')
					echo "max conn is $max_conn"
					printf "allowed are $allows"
					if [[ -n "$allows" ]]; then
						for ip in $allows; do
							whited=$(echo "$connected_rules" | grep "$ip.*:$port" | sort -r)
							if [[ -z "$whited" ]]; then
								#echo "has whited!"
								sudo iptables -I INPUT -p tcp --dport "$port" -s "$ip" -j ACCEPT
							fi
						done
					fi
					has_reject=$(echo "$connected_rules" | grep "DROP.*:$port" | sort -r)
					if [[ -z "$has_reject" ]]; then
						#echo "has reject!"
						sudo iptables -A INPUT -p tcp --dport "$port" -j DROP #--reject-with tcp-reset
					fi	
				else
					rules_number=$(echo "$connected_rules" | awk '{print $1}' | sort -r)
					for rule_number in $rules_number; do
						sudo iptables -D INPUT "$rule_number"
					done
				fi
			done
			sleep 20
			reject_rules=$(sudo iptables -L INPUT -n --line-numbers | grep "DROP" | awk '{print $1}')
			if [[ ! -z "$reject_rules" ]]; then
				for rejectRule in $reject_rules; do
					sudo iptables -D INPUT "$rejectRule"
				done
			fi
			sleep 40
		done;
	elif [[ "$type" == "speed" ]]; then
		logger "running by SPEED limit!"
		echo "running by SPEED limit!"
		tc qdisc add dev "$interface" handle 1: root htb default 30
		tc qdisc add dev eth0 handle ffff: ingress
		echo "$interface"
		while true; do
			connections=$(netstat -tn | awk '{split($4, lparts, ":"); split($5, rparts, ":"); ip=rparts[1] ":" lparts[2]; if (!seen[ip]++) print ip;}')
			rules=$(sudo iptables -L INPUT -n --line-numbers)
			mangles=$(iptables -t mangle -L OUTPUT -n --line-numbers)
			#mangles=$(iptables -t mangle -n --line-numbers)
			ports_setting=$(get_ports | grep -v '^$')
			tc_classes=$(tc class show dev "$interface")
			tc_filters=$(tc filter show dev "$interface")
			tc_filter_ing=$(tc filter show dev eth0 parent ffff:)
			for port in "${ports[@]}"; do
				echo "inside loop"
				port_id=${port: -4}
				echo "$port_id"
				connected_ips=$(echo "$connections" | grep ":$port" | grep -v '^$' | sort -n)
				#connected_ips=$(grep ":$port" | grep -v '^$' | sort -n <<< "$connections")
				echo "$connected_ips"
				#connected_rules=$(echo "$rules" | grep ":$port")
				connected_rules=$(grep ":$port" <<< "$rules")
				#connected_mangles=$(echo "$mangles" | grep ":$port")
				connected_mangles=$(grep ":$port" <<< "$mangles")
				echo "connected mangles: $connected_mangles"
				#count=$(echo "$connected_ips" | wc -l)
				count=$(wc -l <<< "$connected_ips")
				echo "count $count"
				#port_setting=$(echo "$ports_setting" | grep "$port")
				port_setting=$(grep "$port" <<< "$ports_setting")
				echo "current setting: $port_setting";
				port_ignore=$(echo "$port_setting" | awk -F ',' '{print $4}' | tr -d '[:space:]')
				#port_ignore=$(awk -F ',' '{print $4}' | tr -d '[:space:]' <<< "$port_setting")
				echo "ignore: $port_ignore"
				port_speed=$(echo "$port_setting" | awk -F ',' '{print $3}'| tr -d '[:space:]')
				#port_speed=$(awk -F ',' '{print $3}'| tr -d '[:space:]' <<< "$port_setting")
				#echo "speed: $port_speed"
				port_limit=$(echo "$port_setting" | awk -F ',' '{print $2}'| tr -d '[:space:]')
				#port_limit=$(awk -F ',' '{print $2}'| tr -d '[:space:]' <<< "$port_setting")
				#echo "limit: $port_limit"
				port_class=$(echo "$tc_classes" | grep ":$port_id" | grep -v '^$')
				#port_class=$(grep ":$port_id" | grep -v '^$' <<< "$tc_classes")
				port_filter=$(echo "$tc_filters" | grep ":$port_id" | grep -v '^$')
				#port_filter=$(grep ":$port_id" | grep -v '^$' <<< "$tc_filters")
				port_filter_ing=$(echo "$tc_filter_ing" | grep ":$port_id" | grep -v '^$')
				#port_filter_ing=$(grep ":$port_id" | grep -v '^$' <<< "$tc_filter_ing" )
				echo "FILTERS: $port_filter"
				current_speed=$(awk '{print $10}' <<< "$port_class")
				if [ -n "$port_limit" ]; then
					max_conn="$port_limit"
				fi
				if [ -n "$port_speed" ]; then
					speed_limit="$port_speed"
					#echo "has db speed!"
				fi
				if [ "$count" -gt "$max_conn" ] && { [ -z "$port_ignore" ] || [ "$port_ignore" -ne 1 ]; }; then
					total_conn=$((max_conn + speed_limit_th))
					allows=$(echo "$connected_ips" | head -n "$total_conn" | awk -F: '{print $1}')
					#allows=$(head -n "$total_conn" | awk -F: '{print $1}' <<< "$connected_ips")
					echo "allowed ips '$allows'"
					if [[ -n "$allows" ]]; then
						for ip in $allows; do
							whited=$(echo "$connected_rules" | grep "$ip.*:$port" | sort -r)
							#whited=$(grep "$ip.*:$port" | sort -r <<< "$connected_rules")
							if [[ -z "$whited" ]]; then
								#echo "has whited!"
								sudo iptables -I INPUT -p tcp --dport "$port" -s "$ip" -j ACCEPT
							fi
						done
						limit_start=$((max_conn+1))
						limit_end=$((limit_start+speed_limit_th-1))
						#limits=$(echo "$allows" | sed -n "${limit_start},${limit_end}p")
						limits=$(sed -n "${limit_start},${limit_end}p" <<< "$allows")
						echo "limited ips '$limits'"
						if [[ -n "$limits" ]]; then
							echo "$port_id"
							if [[ -z "$port_class" ]]; then
								echo "adding Class for port $port"
								tc class add dev "$interface" parent 1: classid 1:"$port_id" htb rate "$speed_limit"
							else
								echo "port ($port) current tc class: $port_class"
							fi
							if [[ -z "$port_filter" ]]; then
								echo "adding filter for port $port"
								tc filter add dev "$interface" parent 1: prio 1 protocol ip handle 1 fw classid 1:"$port_id"
							fi
							if [ "${current_speed,,}" != "${speed_limit,,}" ]; then
								# Change the speed for the class
								tc class change dev "$interface" classid 1:"$port_id" htb rate "$speed_limit"
							fi
							IFS=$'\n'
							echo "setting speed limit for ip '$ip' with port '$port' to '$speed_limit'"
							for ip in $(echo "$limits"); do
								ip_mangle=$(grep "$ip" <<< "$connected_mangles")
								echo "ip mangle: $ip_mangle"
								if [[ -z "$ip_mangle" ]]; then
									echo "setting download speed for ip $ip"
									iptables -t mangle -A OUTPUT -p tcp -d "$ip" --sport "$port" -j MARK --set-mark 1
									#tc filter del dev "$interface" protocol ip prio 1 u32 match ip dst "$ip" match ip sport "$port" 0xffff flowid 1:"$port_id"
									#tc filter add dev "$interface" protocol ip prio 1 u32 match ip dst "$ip" match ip sport "$port" 0xffff flowid 1:"$port_id"
								fi
								
								if [[ -z "$port_filter_ing" ]]; then
									#echo "setting upload speed for ip $ip"
									#tc filter del dev eth0 parent ffff: protocol ip prio 1 u32 match ip src "$ip" match ip dport "$port" 0xffff police rate "$speed_limit" burst 14k drop flowid :"$port_id"
									tc filter add dev eth0 parent ffff: protocol ip prio 1 u32 match ip src "$ip" match ip dport "$port" 0xffff police rate "$speed_limit" burst 14k drop flowid :"$port_id"
								fi
								
								if [ "${current_speed,,}" != "${speed_limit,,}" ]; then
									# Change the speed for the upload filter
									echo "changing upload speed current: $current_speed and speed limit: $speed_limit"
									tc filter del dev eth0 parent ffff: protocol ip prio 1 u32 match ip src "$ip" match ip dport "$port" 0xffff police rate "$speed_limit" burst 14k drop flowid :"$port_id"
									tc filter add dev eth0 parent ffff: protocol ip prio 1 u32 match ip src "$ip" match ip dport "$port" 0xffff police rate "$speed_limit" burst 14k drop flowid :"$port_id"
								fi
							done
						fi
					fi
					has_reject=$(echo "$connected_rules" | grep "DROP" | sort -r)
					#has_reject=$(grep "DROP" | sort -r <<< "$connected_rules")
					#if [[ -z "$has_reject" && "$count" -gt "$total_conn" ]]; then 
					if [[ -z "$has_reject" ]]; then
						#echo "has reject!"
						sudo iptables -A INPUT -p tcp --dport "$port" -j DROP #--reject-with tcp-reset
					fi
				else
					echo "removing $port"
					rules_number=$(echo "$connected_rules" | awk '{print $1}' | sort -r)
					mangles_number=$(echo "$connected_mangles" | awk '{print $1}' | sort -r)
					for rule_number in $rules_number; do
						sudo iptables -D INPUT "$rule_number"
					done
					for mangle_number in $mangles_number; do
						sudo iptables -D OUTPUT "$mangle_number"
					done
					if [[ -n "$port_class" ]]; then
						tc class del dev "$interface" classid "1:$port_id"
					fi
					if [[ -n "$port_filter" ]]; then
						tc filter del dev "$interface"
						tc filter del dev eth0 parent ffff: protocol ip prio 1 u32 match ip src "$ip" match ip dport "$port" 0xffff police rate "$speed_limit" burst 14k drop flowid :"$port_id"
					fi
				fi
			done
			sleep 10
		done;
	else
		logger "Wrong limit type in config file: $config_file"
		echo "Wrong limit type in config file: $config_file"
		exit 1
	fi
else
	logger "Configuration file not found: $config_file"
	echo "Configuration file not found: $config_file"
	exit 1
fi