#!/bin/bash
#
# Copyright (C) 2015 OpenWrt-dist
# Copyright (C) 2015 Jian Chang <aa65535@live.com>
#
# This is free software, licensed under the GNU General Public License v3.
# See /LICENSE for more information.
#

TAG="SS_SPEC"          # iptables tag
IPT="iptables -t nat"  # alias of iptables
FWI=$(uci get firewall.shadowsocks.path 2>/dev/null)  # firewall include file
IP_REGEX="^([0-9]{1,3}\.){3}[0-9]{1,3}"		      # used to check if input is a valid IP

usage() {
	cat <<-EOF

    Copyright (C) 2015 OpenWrt-dist
    Copyright (C) 2015 Jian Chang <aa65535@live.com>

    Usage: ss-nat [options]

    Valid options are:

        -s <server_ip>          hostname (requires dig) or ip address of shadowsocks remote server
        -l <local_port>         port number of shadowsocks local server
        -S <server_ip>          hostname (requires dig) or ip address of shadowsocks remote UDP server
        -L <local_port>         port number of shadowsocks local UDP server
        -i <ip_list_file>       a file content is bypassed ip list
        -I <interface>          lan interface of nat, default: eth0
        -a <lan_ips>            lan ip of access control, need a prefix to
        define access control mode
        -b <wan_ips>            wan ip of will be bypassed
        -w <wan_ips>            wan ip of will be forwarded
        -e <extra_options>      extra options for iptables
        -o                      apply the rules to the OUTPUT chain
        -u                      enable udprelay mode, TPROXY is required
        -U                      enable udprelay mode, using different IP
        and ports for TCP and UDP
        -f                      flush the rules
        -h                      show this help message and exit

    This is free software, licensed under the GNU General Public License v3.
    See /LICENSE for more information.

EOF
	exit $1
}

loger() {
	# 1.alert 2.crit 3.err 4.warn 5.notice 6.info 7.debug
	logger -st ss-nat[$$] -p$1 $2
}

flush_r() {
	iptables-save -c | grep -v "$TAG" | iptables-restore -c
	ip rule del fwmark 0x01/0x01 table 100 2>/dev/null
	ip route del local 0.0.0.0/0 dev lo table 100 2>/dev/null
	ipset -X ss_spec_lan_ac 2>/dev/null
	ipset -X ss_spec_wan_ac 2>/dev/null
	[ -n "$FWI" ] && echo '#!/bin/sh' >$FWI
	return 0
}

ipset_r() {
	ipset -! -R <<-EOF || return 1
		create ss_spec_wan_ac hash:net
		$(gen_iplist | sed "/^\s*$/d" | sed -e "s/^/add ss_spec_wan_ac /")
		$(for ip in $WAN_FW_IP; do echo "add ss_spec_wan_ac $ip nomatch"; done)
EOF
	$IPT -N SS_SPEC_WAN_AC && \
	$IPT -A SS_SPEC_WAN_AC -m set --match-set ss_spec_wan_ac dst -j RETURN && \
	$IPT -A SS_SPEC_WAN_AC -j SS_SPEC_WAN_FW
	return $?
}

fw_rule() {
	$IPT -N SS_SPEC_WAN_FW && \
	$IPT -A SS_SPEC_WAN_FW -p tcp \
		-j REDIRECT --to-ports $local_port 2>/dev/null || {
		loger 3 "Can't redirect, please check the iptables."
		exit 1
	}
	return $?
}

ac_rule() {
	if [ -n "$LAN_AC_IP" ]; then
		case "${LAN_AC_IP:0:1}" in
			w|W)
				MATCH_SET="-m set --match-set ss_spec_lan_ac src"
				;;
			b|B)
				MATCH_SET="-m set ! --match-set ss_spec_lan_ac src"
				;;
			*)
				loger 3 "Illegal argument \`-a $LAN_AC_IP\`."
				return 2
				;;
		esac
	fi
	IFNAME=${IFNAME:-eth0}
	ipset -! -R <<-EOF || return 1
		create ss_spec_lan_ac hash:net
		$(for ip in ${LAN_AC_IP:1}; do echo "add ss_spec_lan_ac $ip"; done)
EOF
	$IPT -I PREROUTING 1 ${IFNAME:+-i $IFNAME} -p tcp $EXT_ARGS $MATCH_SET \
		-j SS_SPEC_WAN_AC
	if [ "$OUTPUT" = 1 ]; then
		$IPT -I OUTPUT 1 -p tcp $EXT_ARGS -j SS_SPEC_WAN_AC
	fi
	return $?
}

tp_rule() {
	lsmod | grep -q TPROXY || return 0
	[ -n "$TPROXY" ] || return 0
	ip rule add fwmark 0x01/0x01 table 100
	ip route add local 0.0.0.0/0 dev lo table 100
	local ipt="iptables -t mangle"
	$ipt -N SS_SPEC_TPROXY
	$ipt -A SS_SPEC_TPROXY -p udp -m set ! --match-set ss_spec_wan_ac dst \
		-j TPROXY --on-port "$LOCAL_PORT" --tproxy-mark 0x01/0x01
	$ipt -I PREROUTING 1 ${IFNAME:+-i $IFNAME} -p udp $EXT_ARGS $MATCH_SET \
		-j SS_SPEC_TPROXY
	return $?
}

get_wan_ip() {
	cat <<-EOF | grep -E $IP_REGEX
		$server
		$SERVER
		$WAN_BP_IP
EOF
}

gen_iplist() {
	cat <<-EOF
		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.88.99.0/24
		192.168.0.0/16
		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
		$(get_wan_ip)
		$(cat ${IGNORE_LIST:=/dev/null} 2>/dev/null)
EOF
}

gen_include() {
	[ -n "$FWI" ] || return 0
	cat <<-EOF >>$FWI
	iptables-restore -n <<-EOT
	$(iptables-save | grep -E "$TAG|^\*|^COMMIT" |\
			sed -e "s/^-A \(OUTPUT\|PREROUTING\)/-I \1 1/")
	EOT
EOF
	return $?
}

while getopts ":s:l:S:L:i:I:e:a:b:w:ouUfh" arg; do
	case "$arg" in
		s)
            if [[ $OPTARG =~ $IP_REGEX ]]
            then
                server=$OPTARG
            else
                server=$(ping -4 -q -c 1 -s 0 -W 1 -w 1 $OPTARG| head -n 1 | sed -n 's/[^(]*(\([^)]*\)).*/\1/p')
            fi
			;;
		l)
			local_port=$OPTARG
			;;
		S)
            if [[ $OPTARG =~ $IP_REGEX ]]
            then
                SERVER=$OPTARG
            else
                SERVER=$(ping -4 -q -c 1 -s 0 -W 1 -w 1 $OPTARG| head -n 1 | sed -n 's/[^(]*(\([^)]*\)).*/\1/p')
            fi
			;;
		L)
			LOCAL_PORT=$OPTARG
			;;
		i)
			IGNORE_LIST=$OPTARG
			;;
		I)
			IFNAME=$OPTARG
			;;
		e)
			EXT_ARGS=$OPTARG
			;;
		a)
			LAN_AC_IP=$OPTARG
			;;
		b)
			WAN_BP_IP=$(for ip in $OPTARG; do echo $ip; done)
			;;
		w)
			WAN_FW_IP=$OPTARG
			;;
		o)
			OUTPUT=1
			;;
		u)
			TPROXY=1
			;;
		U)
			TPROXY=2
			;;
		f)
			flush_r
			exit 0
			;;
		h)
			usage 0
			;;
	esac
done

if [ -z "$server" -o -z "$local_port" ]; then
	usage 2
fi

if [ "$TPROXY" = 1 ]; then
	SERVER=$server
	LOCAL_PORT=$local_port
elif [ "$TPROXY" = 2 ]; then
	: ${SERVER:?"You must assign an ip for the udp relay server."}
	: ${LOCAL_PORT:?"You must assign a port for the udp relay server."}
fi

flush_r && fw_rule && ipset_r && ac_rule && tp_rule && gen_include
[ "$?" = 0 ] || loger 3 "Start failed!"
exit $?
