#!/bin/sh -e
#
# Copyright (C) 2017 Yousong Zhou <yszhou4tech@gmail.com>
# Copyright (C) 2018-2021 Ycarus (Yannick Chabanois) <ycarus@zugaina.org> for OpenMPTCProuter
#
# The design idea was derived from ss-rules by Jian Chang <aa65535@live.com>
#
# This is free software, licensed under the GNU General Public License v3.
# See /LICENSE for more information.
#

if [ -f /usr/sbin/iptables-legacy ]; then
	IPTABLES="/usr/sbin/iptables-legacy"
	IPTABLESRESTORE="/usr/sbin/iptables-legacy-restore"
	IPTABLESSAVE="/usr/sbin/iptables-legacy-save"
else
	IPTABLES="/usr/sbin/iptables"
	IPTABLESRESTORE="/usr/sbin/iptables-restore"
	IPTABLESSAVE="/usr/sbin/iptables-save"
fi



v2r_rules_usage() {
	cat >&2 <<EOF
Usage: v2ray-rules [options]

	-h, --help      Show this help message then exit
	-f, --flush     Flush rules, ipset then exit
	-l <port>       Local port number of ss-redir with TCP mode
	-L <port>       Local port number of ss-redir with UDP mode
	-s <ips>        List of ip addresses of remote shadowsocks server
	--ifnames       Only apply rules on packets from these ifnames
	--src-bypass <ips|cidr>
	--src-forward <ips|cidr>
	--src-checkdst <ips|cidr>
	--src-default <bypass|forward|checkdst>
	                Packets will have their src ip checked in order against
	                bypass, forward, checkdst list and will bypass, forward
	                through, or continue to have their dst ip checked
	                respectively on the first match.  Otherwise, --src-default
	                decide the default action
	--dst-bypass <ips|cidr>
	--dst-forward <ips|cidr>
	--dst-bypass-file <file>
	--dst-forward-file <file>
	--dst-default <bypass|forward>
	                Same as with their --src-xx equivalent
	--dst-forward-recentrst
	                Forward those packets whose destinations have recently
	                sent to us multiple tcp-rst packets
	--local-default <bypass|forward|checkdst>
	                Default action for local out TCP traffic

The following ipsets will be created by ss-rules.  They are also intended to be
populated by other programs like dnsmasq with ipset support

	ss_rules_src_bypass
	ss_rules_src_forward
	ss_rules_src_checkdst
	ss_rules_dst_bypass
	ss_rules_dst_bypass_all
	ss_rules_dst_forward
EOF
}

o_dst_bypass_="
	0.0.0.0/8
	10.0.0.0/8
	100.64.0.0/10
	127.0.0.0/8
	169.254.0.0/16
	172.16.0.0/12
	192.0.0.0/24
	192.0.2.0/24
	192.31.196.0/24
	192.52.193.0/24
	192.88.99.0/24
	192.168.0.0/16
	192.175.48.0/24
	198.18.0.0/15
	198.51.100.0/24
	203.0.113.0/24
	224.0.0.0/4
	240.0.0.0/4
	255.255.255.255
"
o_src_default=bypass
o_dst_default=bypass
o_local_default=bypass

__errmsg() {
	echo "v2ray-rules: $*" >&2
}

v2r_rules_parse_args() {
	while [ "$#" -gt 0 ]; do
		case "$1" in
			-h|--help) v2r_rules_usage; exit 0;;
			-f|--flush) v2r_rules_flush; exit 0;;
			-l) o_redir_tcp_port="$2"; shift 2;;
			-L) o_redir_udp_port="$2"; shift 2;;
			-s) o_remote_servers="$2"; shift 2;;
			--ifnames) o_ifnames="$2"; shift 2;;
			--ipt-extra) o_ipt_extra="$2"; shift 2;;
			--src-default) o_src_default="$2"; shift 2;;
			--dst-default) o_dst_default="$2"; shift 2;;
			--local-default) o_local_default="$2"; shift 2;;
			--src-bypass) o_src_bypass="$2"; shift 2;;
			--src-forward) o_src_forward="$2"; shift 2;;
			--src-checkdst) o_src_checkdst="$2"; shift 2;;
			--dst-bypass) o_dst_bypass="$2"; shift 2;;
			--dst-bypass_all) o_dst_bypass_all="$2"; shift 2;;
			--dst-forward) o_dst_forward="$2"; shift 2;;
			--dst-forward-recentrst) o_dst_forward_recentrst=1; shift 1;;
			--dst-bypass-file) o_dst_bypass_file="$2"; shift 2;;
			--dst-forward-file) o_dst_forward_file="$2"; shift 2;;
			--rule-name) rule="$2"; shift 2;;
			*) __errmsg "unknown option $1"; return 1;;
		esac
	done

	if [ -z "$o_redir_tcp_port" -a -z "$o_redir_udp_port" ]; then
		__errmsg "Requires at least -l or -L option"
		return 1
	fi
	if [ -n "$o_dst_forward_recentrst" ] && ! $IPTABLES -w -m recent -h >/dev/null; then
		__errmsg "Please install iptables-mod-conntrack-extra with opkg"
		return 1
	fi
	o_remote_servers="$(for s in $o_remote_servers; do resolveip -4 "$s"; done)"
}

v2r_rules_flush() {
	local setname

	$IPTABLESSAVE --counters 2>/dev/null | grep -v v2r_ | $IPTABLESRESTORE --counters
	while ip rule del fwmark 1 lookup 100 2>/dev/null; do true; done
	ip route flush table 100 || true
	for setname in $(ipset -n list | grep "ssr_${rule}"); do
		ipset destroy "$setname" 2>/dev/null || true
	done
}

v2r_rules_ipset_init() {
	ipset --exist restore <<-EOF
		create ssr_${rule}_src_bypass hash:net hashsize 64
		create ssr_${rule}_src_forward hash:net hashsize 64
		create ssr_${rule}_src_checkdst hash:net hashsize 64
		create ss_rules_dst_bypass_all hash:net hashsize 64
		create ssr_${rule}_dst_bypass hash:net hashsize 64
		create ssr_${rule}_dst_bypass_ hash:net hashsize 64
		create ssr_${rule}_dst_forward hash:net hashsize 64
		create ss_rules_dst_forward_recentrst_ hash:ip hashsize 64 timeout 3600
		$(v2r_rules_ipset_mkadd ssr_${rule}_dst_bypass_ "$o_dst_bypass_ $o_remote_servers")
		$(v2r_rules_ipset_mkadd ss_rules_dst_bypass_all "$o_dst_bypass_all")
		$(v2r_rules_ipset_mkadd ssr_${rule}_dst_bypass "$o_dst_bypass $(cat "$o_dst_bypass_file" 2>/dev/null | grep -o '[0-9]\{1,3\}\.[0-9]\{1,3\}\.[0-9]\{1,3\}\.[0-9]\{1,3\}')")
		$(v2r_rules_ipset_mkadd ssr_${rule}_src_bypass "$o_src_bypass")
		$(v2r_rules_ipset_mkadd ssr_${rule}_src_forward "$o_src_forward")
		$(v2r_rules_ipset_mkadd ssr_${rule}_src_checkdst "$o_src_checkdst")
		$(v2r_rules_ipset_mkadd ssr_${rule}_dst_forward "$o_dst_forward $(cat "$o_dst_forward_file" 2>/dev/null | grep -o '[0-9]\{1,3\}\.[0-9]\{1,3\}\.[0-9]\{1,3\}\.[0-9]\{1,3\}')")
	EOF
}

v2r_rules_ipset_mkadd() {
	local setname="$1"; shift
	local i

	for i in $*; do
		echo "add $setname $i"
	done
}

v2r_rules_iptchains_init() {
	v2r_rules_iptchains_init_mark
	v2r_rules_iptchains_init_tcp
	v2r_rules_iptchains_init_udp
}

v2r_rules_iptchains_init_mark() {
	if [ "$($IPTABLES -w -t mangle -L PREROUTING | grep ss_rules_dst_bypass_all)" = "" ]; then
		$IPTABLESRESTORE --noflush <<-EOF
			*mangle
			-A PREROUTING -m set --match-set ss_rules_dst_bypass_all dst -j MARK --set-mark 0x539
			COMMIT
		EOF
	fi
}

v2r_rules_iptchains_init_tcp() {
	local local_target

	[ -n "$o_redir_tcp_port" ] || return 0

	v2r_rules_iptchains_init_ nat tcp

	case "$o_local_default" in
		checkdst) local_target=v2r_${rule}_dst ;;
		forward) local_target=v2r_${rule}_forward ;;
		bypass|*) return 0;;
	esac

	$IPTABLESRESTORE --noflush <<-EOF
		*nat
		:v2r_${rule}_local_out -
		-I OUTPUT 1 -p tcp -j v2r_${rule}_local_out
		-A v2r_${rule}_local_out -m set --match-set ssr_${rule}_dst_bypass dst -j RETURN
		-A v2r_${rule}_local_out -m set --match-set ss_rules_dst_bypass_all dst -j RETURN
		-A v2r_${rule}_local_out -m set --match-set ssr_${rule}_dst_bypass_ dst -j RETURN
		-A v2r_${rule}_local_out -m mark --mark 0x539 -j RETURN
		-A v2r_${rule}_local_out -p tcp $o_ipt_extra -j $local_target -m comment --comment "local_default: $o_local_default"
		COMMIT
	EOF
}

v2r_rules_iptchains_init_udp() {
	[ -n "$o_redir_udp_port" ] || return 0
	v2r_rules_iptchains_init_ mangle udp
}

v2r_rules_iptchains_init_() {
	local table="$1"
	local proto="$2"
	local forward_rules
	local src_default_target dst_default_target
	local recentrst_mangle_rules recentrst_addset_rules

	case "$proto" in
		tcp)
			forward_rules="-A v2r_${rule}_forward -p tcp -j REDIRECT --to-ports $o_redir_tcp_port"
			if [ -n "$o_dst_forward_recentrst" ]; then
				recentrst_mangle_rules="
					*mangle
					-I PREROUTING 1 -p tcp -m tcp --tcp-flags RST RST -m recent --name v2r_recentrst --set --rsource
					COMMIT
				"
				recentrst_addset_rules="
					-A v2r_${rule}_dst -m recent --name v2r_recentrst --rcheck --rdest --seconds 3 --hitcount 3 -j SET --add-set ss_rules_dst_forward_recentrst_ dst --exist
					-A v2r_${rule}_dst -m set --match-set ss_rules_dst_forward_recentrst_ dst -j v2r_${rule}_forward
				"
			fi
			;;
		udp)
			ip rule add fwmark 1 lookup 100 || true
			ip route add local default dev lo table 100 || true
			forward_rules="-A v2r_${rule}_forward -p udp -j TPROXY --on-port "$o_redir_udp_port" --tproxy-mark 0x01/0x01"
			;;
	esac
	case "$o_src_default" in
		forward) src_default_target=v2r_${rule}_forward ;;
		checkdst) src_default_target=v2r_${rule}_dst ;;
		bypass|*) src_default_target=RETURN ;;
	esac
	case "$o_dst_default" in
		forward) dst_default_target=v2r_${rule}_forward ;;
		bypass|*) dst_default_target=RETURN ;;
	esac
	sed -e '/^\s*$/d' -e 's/^\s\+//' <<-EOF | $IPTABLESRESTORE --noflush
		*$table
		:v2r_${rule}_pre_src -
		:v2r_${rule}_src -
		:v2r_${rule}_dst -
		:v2r_${rule}_forward -
		$(v2r_rules_iptchains_mkprerules "$proto")
		-A v2r_${rule}_pre_src -m set --match-set ssr_${rule}_dst_bypass_ dst -j RETURN
		-A v2r_${rule}_pre_src -m set --match-set ss_rules_dst_bypass_all dst -j MARK --set-mark 0x539
		-A v2r_${rule}_pre_src -m set --match-set ss_rules_dst_bypass_all dst -j RETURN
		-A v2r_${rule}_pre_src -m set --match-set ssr_${rule}_dst_bypass dst -j RETURN
		-A v2r_${rule}_pre_src -m mark --mark 0x539 -j RETURN
		-A v2r_${rule}_dst -m set --match-set ss_rules_dst_bypass_all dst -j RETURN
		-A v2r_${rule}_dst -m set --match-set ssr_${rule}_dst_bypass dst -j RETURN
		-A v2r_${rule}_pre_src -p $proto $o_ipt_extra -j v2r_${rule}_src
		-A v2r_${rule}_src -m set --match-set ssr_${rule}_src_bypass src -j RETURN
		-A v2r_${rule}_src -m set --match-set ssr_${rule}_src_forward src -j v2r_${rule}_forward
		-A v2r_${rule}_src -m set --match-set ssr_${rule}_src_checkdst src -j v2r_${rule}_dst
		-A v2r_${rule}_src -j $src_default_target -m comment --comment "src_default: $o_src_default"
		-A v2r_${rule}_dst -m set --match-set ssr_${rule}_dst_forward dst -j v2r_${rule}_forward
		$recentrst_addset_rules
		-A v2r_${rule}_dst -j $dst_default_target -m comment --comment "dst_default: $o_dst_default"
		$forward_rules
		COMMIT
		$recentrst_mangle_rules
	EOF
}

v2r_rules_iptchains_mkprerules() {
	local proto="$1"

	if [ -z "$o_ifnames" ]; then
		echo "-A PREROUTING -p $proto -j v2r_${rule}_pre_src"
	else
		echo $o_ifnames \
			| tr ' ' '\n' \
			| sed "s/.*/-I PREROUTING 1 -i \\0 -p $proto -j v2r_${rule}_pre_src/"
	fi
}

v2r_rules_fw_drop() {
	fw3 -4 print 2>/dev/null | awk '/iptables/&&/zone_lan_forward/&&/tcp/&&/-t filter/&&/-j reject/ {for(i=6; i<=NF; i++) { printf "%s ",$i } print "\n" }' |
	while IFS=$"\n" read -r c; do
		fwrule=$(echo "$c" | sed 's/reject/REDIRECT --to-ports 65535/')
		if [ -n "$fwrule" ] && [ -z "$($IPTABLESSAVE 2>/dev/null | grep zone_lan_prerouting | grep '${fwrule}')" ]; then
			eval "$IPTABLES -w -t nat -A zone_lan_prerouting ${fwrule} 2>&1 >/dev/null"
		fi
	done
	fw3 -4 print 2>/dev/null | awk '/iptables/&&/zone_lan_forward/&&/tcp/&&/-t filter/&&/-j drop/ {for(i=6; i<=NF; i++) { printf "%s ",$i } print "\n" }' |
	while IFS=$"\n" read -r c; do
		fwrule=$(echo "$c" | sed 's/drop/REDIRECT --to-ports 65535/')
		if [ -n "$fwrule" ] && [ -z "$($IPTABLESSAVE 2>/dev/null | grep zone_lan_prerouting | grep '${fwrule}')" ]; then
			eval "$IPTABLES -t nat -A zone_lan_prerouting ${fwrule} 2>&1 >/dev/null"
		fi
	done
}

v2r_rules_parse_args "$@"
#v2r_rules_flush
v2r_rules_ipset_init
v2r_rules_iptchains_init
v2r_rules_fw_drop