#!/bin/sh
#
# Copyright (C) 2018-2023 Ruilin Peng (Nick) <pymumu@gmail.com>.
#
# smartdns is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# smartdns is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

SMARTDNS_BIN=/opt/usr/sbin/smartdns
SMARTDNS_CONF=/opt/etc/smartdns/smartdns.conf
DNSMASQ_CONF="/etc/dnsmasq.conf /var/etc/dnsmasq.conf /etc/storage/dnsmasq/dnsmasq.conf"
SMARTDNS_PID=/var/run/smartdns.pid
SMARTDNS_PORT=535
SMARTDNS_OPT=/opt/etc/smartdns/smartdns-opt.conf
# workmode 
# DO NOT CHANGE THIS, CHANGE MODE IN smartdns-opt.conf
# 0: run as port only
# 1: redirect port
# 2: replace 
SMARTDNS_WORKMODE="1"

if [ -f "$SMARTDNS_OPT" ]; then
. "$SMARTDNS_OPT"
fi


set_iptable()
{
	local redirect_tcp

	redirect_tcp=0

	grep ^bind-tcp $SMARTDNS_CONF > /dev/null 2>&1
	if [ $? -eq 0 ]; then
		redirect_tcp=1;
	fi

	IPS="$(ifconfig | grep "inet addr" | grep -v ":127" | grep "Bcast" | awk '{print $2}' | awk -F: '{print $2}')"
	for IP in $IPS
	do
		if [ $redirect_tcp -eq 1 ]; then
			iptables -t nat -A PREROUTING -p tcp -d "$IP" --dport 53 -j REDIRECT --to-ports "$SMARTDNS_PORT" > /dev/null 2>&1
		fi
		iptables -t nat -A PREROUTING -p udp -d "$IP" --dport 53 -j REDIRECT --to-ports "$SMARTDNS_PORT" > /dev/null 2>&1
	done

}

clear_iptable()
{
	IPS="$(ifconfig | grep "inet addr" | grep -v ":127" | grep "Bcast" | awk '{print $2}' | awk -F: '{print $2}')"
	for IP in $IPS
	do
		iptables -t nat -D PREROUTING -p tcp -d "$IP" --dport 53 -j REDIRECT --to-ports "$SMARTDNS_PORT" > /dev/null 2>&1
		iptables -t nat -D PREROUTING -p udp -d "$IP" --dport 53 -j REDIRECT --to-ports "$SMARTDNS_PORT" > /dev/null 2>&1
	done
	
}

get_dnsmasq_cmd()
{
	CMD="$(ps 2>/dev/null | grep -e '[a-zA-Z]\{0,2\} \{1,\}dnsmasq' | grep -v grep 2>/dev/null)"
	if [ ! -z "$CMD" ]; then
		return
	fi

	CMD="$(ps 2>/dev/null | grep '/usr/sbin/dnsmasq' | grep -v grep 2>/dev/null)"
	if [ ! -z "$CMD" ]; then
		return
	fi

	CMD="$(ps 2>/dev/null | grep 'dnsmasq' | grep -v grep 2>/dev/null)"
	if [ ! -z "$CMD" ]; then
		return
	fi

	CMD="$(ps ax 2>/dev/null | grep -e '[a-zA-Z]\{0,2\} \{1,\}dnsmasq' | grep -v grep 2>/dev/null)"
	if [ ! -z "$CMD" ]; then
		return
	fi

	CMD="$(ps ax 2>/dev/null | grep /usr/sbin/dnsmasq | grep -v grep 2>/dev/null)"
	if [ ! -z "$CMD" ]; then
		return
	fi

	CMD="$(ps ax 2>/dev/null | grep 'dnsmasq' | grep -v grep 2>/dev/null)"
	if [ ! -z "$CMD" ]; then
		return
	fi	
}

restart_dnsmasq()
{
	local CMD=""

	get_dnsmasq_cmd
	if [ -z "$CMD" ]; then
		echo "cannot find dnsmasq"
		return 1
	fi

	# check multiple dnsmasq
	linecount="$(echo "$CMD" | wc -l)"
	if [ $linecount -eq 1 ]; then
		PID="$(echo "$CMD" | awk '{print $1}')"
	elif [ $linecount -gt 1 ]; then
		PID1="$(echo "$CMD" | awk 'NR==1{print $1}')"
		PID2="$(echo "$CMD" | awk 'NR==2{print $1}')"
		PID2_PPID="$(grep 'PPid:' /proc/$PID2/status | awk '{print $2}' 2>/dev/null)"
		if [ "$PID2_PPID" != "$PID1" ]; then
			echo "find multiple dnsmasq, but not started by the same process"
			return 1
		fi
		PID=$PID1
	else
		echo "find multiple dnsmasq, but not started by the same process"
		return 1
	fi

	if [ ! -d "/proc/$PID" ]; then
		echo "dnsmasq is not running"
		return 1
	fi

	kill -9 "$PID"

	# get dnsmasq command
	CMD="$(echo "$CMD" | head -n 1)"
	DNSMASQ_CMD="$(echo "$CMD" | awk '{for(i=5; i<=NF;i++)printf $i " "}')"

	$DNSMASQ_CMD
}

add_dhcp_options6()
{
	CONF_FILE=$1
	IPS="$(ifconfig | grep "inet addr" | grep -v ":127" | grep "Bcast" | awk '{print $2}' | awk -F: '{print $2}')"
	for IP in $IPS
	do
		DHCP_OPTION="$(grep "dhcp-option=" "$CONF_FILE" | grep "$IP" | head -n 1)"
		if [ -z "$DHCP_OPTION" ]; then
			continue
		fi

		SERVER_TAG="$(echo "$DHCP_OPTION" | awk -F= '{print $2}' | awk -F, '{print $1}')"
		LOCAL_SERVER_IP="$IP"

		grep "dhcp-option *= *$SERVER_TAG, *6 *, *$LOCAL_SERVER_IP" "$CONF_FILE" 1>/dev/null 2>&1
		if [ $? -eq 0 ]; then
			continue
		fi
		
		DHCP_OPTION="dhcp-option=$SERVER_TAG,6,$LOCAL_SERVER_IP"
		echo "$DHCP_OPTION" >> "$CONF_FILE"
		RESTART_DNSMASQ=1
	done

	return 1
}

clear_dhcp_options6()
{
	CONF_FILE=$1
	IPS="$(ifconfig | grep "inet addr" | grep -v ":127" | grep "Bcast" | awk '{print $2}' | awk -F: '{print $2}')"
	for IP in $IPS
	do
		DHCP_OPTION="$(grep "dhcp-option=" "$CONF_FILE" | grep "$IP" | head -n 1)"
		if [ -z "$DHCP_OPTION" ]; then
			continue
		fi

		SERVER_TAG="$(echo "$DHCP_OPTION" | awk -F= '{print $2}' | awk -F, '{print $1}')"
		LOCAL_SERVER_IP="$IP"

		grep "dhcp-option *= *$SERVER_TAG, *6 *, *$LOCAL_SERVER_IP" "$CONF_FILE" 1>/dev/null 2>&1
		if [ $? -ne 0 ]; then
			continue
		fi
		
		sed -i "/^dhcp-option *=$SERVER_TAG,6,/d" "$CONF_FILE"
		RESTART_DNSMASQ=1
	done

	return 1
}

set_dnsmasq_conf()
{
	local LOCAL_SERVER_IP=""
	local SERVER_TAG=""
	local CONF_FILE=$1
	local DHCP_OPTIONS=""

	add_dhcp_options6 "$CONF_FILE"

	grep "^port *=0" "$CONF_FILE" > /dev/null 2>&1
	if [ $? -ne 0 ]; then
		sed -i "/^port *=/d" "$CONF_FILE"
		echo "port=0" >> "$CONF_FILE"
		RESTART_DNSMASQ=1
	fi
}

set_dnsmasq()
{
	local RESTART_DNSMASQ=0

	for conf in $DNSMASQ_CONF
	do
		if [ ! -e "$conf" ]; then
			continue
		fi

		set_dnsmasq_conf "$conf"
	done
	
	if [ $RESTART_DNSMASQ -ne 0 ]; then
		restart_dnsmasq	
	fi
}

clear_dnsmasq_conf()
{
	local LOCAL_SERVER_IP=""
	local SERVER_TAG=""
	local CONF_FILE=$1
	
	clear_dhcp_options6 "$CONF_FILE"

	grep "^port *=" "$CONF_FILE" > /dev/null 2>&1
	if [ $? -eq 0 ]; then
		sed -i "/^port *=/d" "$CONF_FILE"
		RESTART_DNSMASQ=1
	fi
}

clear_dnsmasq()
{
	local RESTART_DNSMASQ=0

	for conf in $DNSMASQ_CONF
	do
		if [ ! -e "$conf" ]; then
			continue
		fi

		clear_dnsmasq_conf "$conf"
	done

	if [ $RESTART_DNSMASQ -ne 0 ]; then
		restart_dnsmasq	
	fi
}

set_smartdns_port()
{
	if [ "$SMARTDNS_WORKMODE" = "0" ]; then
		return 0
	elif [ "$SMARTDNS_WORKMODE" = "1" ]; then
		sed -i "s/^\(bind .*\):53\( .*\)\?$/\1:$SMARTDNS_PORT \2/g" $SMARTDNS_CONF
		sed -i "s/^\(bind-tcp .*\):53\( .*\)\?$/\1:$SMARTDNS_PORT \2/g" $SMARTDNS_CONF
	elif [ "$SMARTDNS_WORKMODE" = "2" ]; then
		sed -i "s/^\(bind .*\):$SMARTDNS_PORT\( .*\)\?$/\1:53 \2/g" $SMARTDNS_CONF
		sed -i "s/^\(bind-tcp .*\):$SMARTDNS_PORT\( .*\)\?$/\1:53 \2/g" $SMARTDNS_CONF
	else
		return 1
	fi	

	return 0
}

set_rule()
{
	if [ "$SMARTDNS_WORKMODE" = "0" ]; then
		return 0
	elif [ "$SMARTDNS_WORKMODE" = "1" ]; then
		set_iptable
		return $?
	elif [ "$SMARTDNS_WORKMODE" = "2" ]; then
		set_dnsmasq
		return $?
	else
		return 1
	fi
}

clear_rule()
{
	if [ "$SMARTDNS_WORKMODE" = "0" ]; then
		return 0
	elif [ "$SMARTDNS_WORKMODE" = "1" ]; then
		clear_iptable
		return $?
	elif [ "$SMARTDNS_WORKMODE" = "2" ]; then
		clear_dnsmasq
		return $?
	else
		return 1
	fi
}

get_tz()
{
	if [ -e "/etc/localtime" ]; then
		return 
	fi
	
	for tzfile in /etc/TZ /var/etc/TZ
	do
		if [ ! -e "$tzfile" ]; then
			continue
		fi
		
		tz="$(cat $tzfile 2>/dev/null)"
	done
	
	if [ -z "$tz" ]; then
		return	
	fi
	
	export TZ=$tz
}

case "$1" in
	start)
	set_rule
	if [ $? -ne 0 ]; then
		exit 1
	fi

	set_smartdns_port
	get_tz
	$SMARTDNS_BIN -c "$SMARTDNS_CONF" -p $SMARTDNS_PID
	if [ $? -ne 0 ]; then
		clear_rule
	fi
	;;
	status)
	pid="$(cat $SMARTDNS_PID |head -n 1 2>/dev/null)"
	if [ -z "$pid" ]; then
		echo "smartdns not running."
		return 0
	fi

	if [ -d "/proc/$pid" ]; then
		echo "smartdns running"
		return 0;
	fi
	echo "smartdns not running."
	return 0;
	;;
	stop)
	clear_rule
	pid="$(cat "$SMARTDNS_PID" | head -n 1 2>/dev/null)"
	if [ -z "$pid" ]; then
		echo "smartdns not running."
		return 0
	fi

	kill -15 "$pid" 2>/dev/null
	SLEEP=$(which usleep 2>/dev/null)
	SLEEPTIME=200000
	if [ -z "$SLEEP" ]; then
		SLEEP="sleep"
		SLEEPTIME=0.2
	fi
	N=300
	while [ $N -gt 0 ]
	do
		pid="$(cat "$SMARTDNS_PID" | head -n 1 2>/dev/null)"
		if [ -z "$pid" ]; then
			return 0
		fi

		if [ ! -d "/proc/$pid" ]; then
			return 0
		fi

		stat="$(cat /proc/${pid}/stat | awk '{print $3}' 2>/dev/null)"
		if [ "$stat" = "Z" ]; then
			$SLEEP $SLEEPTIME
			return 0
		fi

		$SLEEP $SLEEPTIME 2>/dev/null
		N=$((N-1))
	done

	kill -9 "$pid" 2>/dev/null
	;;
	restart)
	$0 stop
	$0 start
	;;
	enable)
	nvram set apps_state_enable=2
	nvram set apps_state_error=0
	nvram set apps_state_install=5
	nvram set apps_state_action=install
	nvram set apps_u2ec_ex=2
	;;
	firewall-start|reload|force-reload|reconfigure)
	$0 restart
	;;
	*)
	;;
esac