#!/usr/bin/env bash
# Exit Codes: 0 - success 
#             1 - script error
#             2 - user error 
#             4 - help

help_and_exit() {
	cat 1>&2 << EOF
${BRIGHT}fw_nc.sh${NOCOLOR}: Firewall Netcat by GI Jack

Open port(s) in the firewall, runs netcat, and then remove the firewall rule.

Script gets port and proto from netcat statement. No special options. 
Just use netcat syntax.

	${BRIGHT}Usage${NOCOLOR}:
	fw_nc.sh <netcat statement>
EOF
	exit 4
}

BRIGHT=$(tput bold)
NOCOLOR=$(tput sgr0) # reset to default colors
BRIGHT_RED=$(tput setaf 1;tput bold)
BRIGHT_YELLOW=$(tput setaf 3;tput bold)
BRIGHT_CYAN=$(tput setaf 6;tput bold)

exit_with_error() {
echo 1>&2 "${BRIGHT}fw_nc.sh: ${BRIGHT_RED}ERROR${NOCOLOR}: ${2}"
exit ${1}
}

message() {
	echo "fw_nc.sh: ${@}"
}

check_sudo() {
	# Check if this script can run sudo correctly. 
	# Uses as_root, see below
	if [ ${UID} -eq 0 ];then
		ROOT_METHOD="uid"
		# TODO: FIX Polkit support
		#elif [[ ! -z $DISPLAY && $(tty) = /dev/pts/* ]];then
		# ROOT_METHOD="pkexec"
	elif [ $(sudo whoami) == "root" ];then
		ROOT_METHOD="sudo"
	else
		exit_with_error 4 "Cannot gain root! This program needs root to work. Exiting..."
	fi
	# One last check
	[ $(as_root whoami) != "root" ] && exit_with_error 4 "Cannot gain root! This program needs root to work. Exiting..."
}


as_root() {
	# Execute a command as root
	case $ROOT_METHOD in
		sudo)
			sudo ${@}
			;;
		polkit)
			pkexec ${@}
			;;
		uid)
			${@}
			;;
	esac
}

check_nc_opts() {
	CHAIN="OUTPUT"
	PROTO="tcp"
	PORT=""
	local parms=""
	while [ ! -z "$1" ];do
		item="${1}"
		case ${item} in
			-l)
			PORT=${2}
			CHAIN="INPUT"
			;;
		-u)
			PROTO="udp"
			;;
		-*)
			shift
			continue
			;;
		*)
			parms+=${item}
			;;
			esac
			shift
		done
		[ -z "${parms}" ] && return
		set ${parms}
		[ -z ${PORT} ] && PORT=${2}
}

remove_rule_and_exit() {
	as_root iptables -D ${CHAIN} -m ${PROTO} -p ${PROTO} --dport ${PORT} -j ACCEPT
	exit
}

main() {
	# Sanity Check
	[ -z $1 ] && help_and_exit
	[ $1 == "--help" ] && help_and_exit
	[ -z $PORT ] && exit_with_error 2 "No such port"
	check_sudo
	as_root iptables -I ${CHAIN} -m ${PROTO} -p ${PROTO} --dport ${PORT} -j ACCEPT
	trap remove_rule_and_exit 1 2 3 9 15
	nc ${@}
	remove_rule_and_exit
}

check_nc_opts "${@}"
main "${@}"