#!/bin/bash

## check for root

NAT_TYPE=
UDP_PORT=
IP=
DELETE=0

OUTER_INTERFACE=eth0

OUTER_IP=`ifconfig ${OUTER_INTERFACE} | grep -e "\binet\b" | awk '{print $2}'`

options=$(getopt -o t:p:i: --long type:,port:,ip:,delete -- "$@")
eval set -- "$options"

while true; do
	case "$1" in
	-t|--type)
		NAT_TYPE=$2
		shift 2
		;;
	-p|--port)
		PORT=$2
		shift 2
		;;
	-i|--ip)
		IP=$2
		shift 2
		;;
	--delete)
		DELETE=1
		shift
		;;
	--)
		shift
		break
		;;
	*)
		echo "invalid option: $1"
		exit -1;
		;;
	esac
done

port=$PORT
ip=$IP

is_number() {
	re="^[1-9][0-9]*$"
	if [[ $1 =~ $re ]]; then
		return 0;
	fi
	return 1;
}

is_valid_ip() {
	re="^[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}$"
	if [[ $1 =~ $re ]]; then
		return 0
	fi
	return 1
}

if ! is_number $port; then
	echo "invalid port: $port"
	exit -1;
fi

if ! is_valid_ip $ip; then
	echo "invalid ip: $ip"
	exit -1;
fi

PUNCHNET_TCP_PORT=18083
## iptables' tcp POSTROUTING chain name
POST_TCP=PUNCHNET-TCP
PRE_UDP=PUNCHNET-PRE-UDP

IPSET_GROUP=allowed_ip_${ip}
add_ipset_group() {
	ipset destroy ${IPSET_GROUP}
	ipset create ${IPSET_GROUP} hash:ip hashsize 4096
}

add_ipset_group

ensure_rule() {
	table_name=$1
	chain_name=$2
	arguments=$3

	## create table
	# iptables -t ${table_name} -N ${chain_name} > /dev/null 2>&1

	## add the rule
	if ! iptables -t ${table_name} -C ${chain_name} ${arguments} > /dev/null 2>&1; then
		echo "executing: iptables -t ${table_name} -A ${chain_name} ${arguments} > /dev/null 2>&1"
		iptables -t ${table_name} -A ${chain_name} ${arguments} > /dev/null 2>&1
	fi
}

clear_rule() {
	table_name=$1
	chain_name=$2
	arguments=$3

	iptables -t ${table_name} -D ${chain_name} ${arguments} > /dev/null 2>&1
}

drop_user_chian() {
	table_name=$1
	chain_name=$2

	iptables -t ${table_name} -F ${chain_name}
	iptables -t ${table_name} -X ${chain_name}
}


ensure_tcp_connection() {
	# create table anyhow
	iptables -t nat -N ${POST_TCP} > /dev/null 2>&1

	ensure_rule "nat" "${POST_TCP}" "-p tcp --dport ${PUNCHNET_TCP_PORT} -j MASQUERADE"
	ensure_rule "nat" "${POST_TCP}" "-j RETURN"
	ensure_rule "nat" "POSTROUTING" "-p tcp -s ${ip} -j ${POST_TCP}"
}

fullcone() {
	ensure_tcp_connection

	# POSTROUTING for udp ports
	ensure_rule "nat" "POSTROUTING" "-o ${OUTER_INTERFACE} -p udp --sport ${port} -j SNAT --to-source ${OUTER_IP}:${port}"

	# prerouting for udp to the very host
	ensure_rule "nat" "PREROUTING" "-p udp -i ${OUTER_INTERFACE} --dport ${port} -j DNAT --to-destination ${ip}:${port}"
}

fullcone_clear() {
	clear_rule "nat" "POSTROUTING" "-o ${OUTER_INTERFACE} -p udp --sport ${port} -j MASQUERADE"
	clear_rule "nat" "PREROUTING" "-i ${OUTER_INTERFACE} -p udp --dport ${port} -j DNAT --to-destination ${ip}"
}

restricted_cone() {
	fullcone
	# ensure_rule "filter" "FORWARD" "-i ${OUTER_INTERFACE} -p udp --dport ${port} -m state --state ESTABLISHED,RELATED -j LOG --log-prefix \\\"accept INPUT: \\\""
	ensure_rule "filter" "FORWARD" "-i ${OUTER_INTERFACE} -p udp --dport ${port} -m state --state ESTABLISHED,RELATED -j SET --add-set ${IPSET_GROUP} src"
	ensure_rule "filter" "FORWARD" "-i ${OUTER_INTERFACE} -p udp --dport ${port} -m state --state ESTABLISHED,RELATED -j ACCEPT"
	# ensure_rule "filter" "FORWARD" "-i ${OUTER_INTERFACE} -p udp --dport ${port} -m state --state ESTABLISHED,RELATED -j ACCEPT"
	# ensure_rule "filter" "FORWARD" "-i ${OUTER_INTERFACE} -p udp --dport ${port} -m state --state NEW -j LOG --log-prefix \"drop INPUT\""
	# ensure_rule "filter" "FORWARD" "-i ${OUTER_INTERFACE} -p udp --dport ${port} -m state --state NEW -m set ! --match-set allowed_ip src -j LOG"
	ensure_rule "filter" "FORWARD" "-i ${OUTER_INTERFACE} -p udp --dport ${port} -m state --state NEW -m set ! --match-set ${IPSET_GROUP} src -j DROP"
	ensure_rule "filter" "FORWARD" "-i ${OUTER_INTERFACE} -p udp --dport ${port} -m state --state NEW -m set --match-set ${IPSET_GROUP} src -j ACCEPT"
}

restricted_cone_clear() {
	fullcone_clear
	clear_rule "filter" "FORWARD" "-i ${OUTER_INTERFACE} -p udp --dport ${port} -m state --state ESTABLISHED,RELATED -j ACCEPT"
	clear_rule "filter" "FORWARD" "-i ${OUTER_INTERFACE} -p udp --dport ${port} -m state --state NEW -j DROP"
}

port_restricted_cone() {
	ensure_tcp_connection
	ensure_rule "nat" "POSTROUTING" "-o ${OUTER_INTERFACE} -p udp --sport ${port} -j MASQUERADE"
}

port_restricted_cone_clear() {
	clear_rule "nat" "POSTROUTING" "-o ${OUTER_INTERFACE} -p udp --sport ${port} -j MASQUERADE"
}

symmetric() {
	ensure_rule "nat" "POSTROUTING" "-o ${OUTER_INTERFACE} -p udp --sport ${port} -j MASQUERADE --random"
}

symmetric_clear() {
	clear_rule "nat" "POSTROUTING" "-o ${OUTER_INTERFACE} -p udp --sport ${port} -j MASQUERADE --random"
}

case $NAT_TYPE in
	"nat1")
		fullcone
		echo "full cone"
		;;
	"nat2")
		restricted_cone
		echo "restricted cone"
		;;
	"nat3")
		port_restricted_cone
		echo "port restricted cone"
		;;
	"nat4")
		symmetric
		echo "symmetric nat"
		;;
	*)
		echo "invalid nat type"
		;;
esac

# iptables -t nat -C POSTROUTING -p tcp -j ${POST_TCP} > /dev/null 2>&1
# 
# 
# iptables -t nat -F ${POST_TCP} > /dev/null 2&>1
# iptables -t nat -N ${POST_TCP}
# iptables -t nat -A POSTROUTING -p tcp -j ${POST_TCP}
# iptables -t nat -A ${POST_TCP} -p tcp --dport ${PUNCHNET_TCP_PORT}
# iptables -t nat -A ${POST_TCP} -j RETURN
# 
# iptables -t nat -D POSTROUTING -p udp -j ${POST_UDP} > /dev/null 2&>1
# iptables -t nat -F ${POST_UDP} > /dev/null 2&>1
# iptables -t nat -N ${POST_UDP}
# 
# iptables -t nat -A POSTROUTING -o eth0 -p tcp --dport 18083 -j MASQUERADE
# iptables -t nat -A POSTROUTING -o eth0 -p udp --sport 7890 -j MASQUERADE
# iptables -t nat -A PREROUTING -i eth0 -p udp --dport 7890 -j DNAT --to-destination 172.17.0.2