#!/bin/bash

function showHelp() {
	echo "mtm-ssh-proxy [args] command cmd_args"
	echo ""
	echo "args :"
	echo "  --help    Show help"
	echo "command :"
	echo "  status"
	echo "  open profile"
	echo "  open port ssh"
	echo "  close listen_port"
	echo "  open-all"
	echo "  close-all"
	echo "cmd_args :"
	echo "  profile   Name of the profile (~/.config/mtm-ssh-proxy/*)"
	echo "  port      Proxy port"
	echo "  ssh       SSH connection (user@server:port)"
}

if [ "${1}" == "" ]; then
	showHelp
	exit 1
fi
case ${1} in
	"--help")
		showHelp
		exit 0
	;;
	"status" | "open-all" | "close-all");;
	"open")
		if [ "${2}" == "" ] || [[ "${2}" == */* ]]; then
			showHelp
			exit 1
		fi
	;;
	"close")
		if [ "${2}" == "" ] || [[ "${2}" == */* ]]; then
			showHelp
			exit 1
		fi
	;;
	*)
		showHelp
		exit 1
	;;
esac

declare -r CONFIG_DIR=~/.config/mtm-ssh-proxy
declare -r TMP_DIR=/tmp/mtm-ssh-proxy-"${USER}/proxy"

read_config() {
	declare -r FILE=${1}
	declare -r KEYNAME=${2}
	declare -r DEFAULT=${3}
	if [ -f "${FILE}" ]; then
		if grep -v "#" "${FILE}" | grep -q "${KEYNAME}="; then
			VALUE=$(grep -v "#" "${FILE}" | grep "^${KEYNAME}=" | sed 's/'"${KEYNAME}"'=//')
			echo "${VALUE}"
			return 0
		fi
	fi
	echo "${DEFAULT}"
}

status() {
	if [ -d "${TMP_DIR}" ] && [ "$(find "${TMP_DIR}" -maxdepth 0 -empty -exec echo 1 \;)" == "" ]; then
		for FILE in "${TMP_DIR}"/*; do
			PROFILE=$(basename "${FILE}")
			PROXY=$(grep 'PROXY=' "${FILE}" | sed 's/PROXY=//')
			echo -e "${PROFILE}\t${PROXY}"
		done
	else
		echo "No proxy open"
	fi
}

open() {
	declare LISTEN_PORT=""
	declare PROXY=""
	if [ "${2}" == "" ]; then
		declare -r FILE=${CONFIG_DIR}/profiles/${1}
		if [ ! -f "${FILE}" ]; then
			echo "Error: Profile not found: ${FILE}" >&2
			exit 1
		fi
		LISTEN_PORT=$(read_config "${FILE}" "PROXY_LOCAL_PORT" 0)
		PROXY=$(read_config "${FILE}" "PROXY" "")
	else
		LISTEN_PORT=${1}
		PROXY=${2}
	fi

	# check if port is in range
	if [[ ! ${LISTEN_PORT} =~ ^[0-9]+$ ]] || [[ ! ${LISTEN_PORT} -lt 65536 ]]; then
		echo "Port is not valid (1024-65535)"
		return 1
	fi

	# check if port is not used
	if [ ! "$(lsof -i -P | grep ":${LISTEN_PORT} (LISTEN)" | awk -F ' ' '{print $2}')" == "" ]; then
		echo "Port already open"
		return 1
	fi

	# parse proxy
	declare PROXY_USER=""
	# declare PROXY_PASSWORD=""
	declare PROXY_SERVER=""
	declare PROXY_SERVER_PORT="22"
	declare PROXY_ADDRESS=""
	if [[ "${PROXY}" == *@* ]]; then
		PROXY_USER="${PROXY%@*}"
		PROXY_SERVER="${PROXY#*@*}"
	else
		PROXY_SERVER="${PROXY}"
	fi
	#if [[ "${PROXY_USER}" == *:* ]]; then
	#  PROXY_PASSWORD="${PROXY_USER#*:*}"
	#  PROXY_USER="${PROXY_USER%:*}"
	#fi
	if [[ "${PROXY_SERVER}" == *:* ]]; then
		PROXY_SERVER_PORT="${PROXY_SERVER#*:*}"
		PROXY_SERVER="${PROXY_SERVER%:*}"
	fi
	PROXY_ADDRESS="${PROXY_SERVER}"
	if [ ! "${PROXY_USER}" == "" ]; then
		PROXY_ADDRESS="${PROXY_USER}@${PROXY_SERVER}"
	fi

	# Results
	echo "Open proxy ${LISTEN_PORT} -> ${PROXY}"
	echo "ssh ${PROXY_ADDRESS} -p ${PROXY_SERVER_PORT}"
	if ! ssh -fND 127.0.0.1:"${LISTEN_PORT}" "${PROXY_ADDRESS}" -p "${PROXY_SERVER_PORT}"; then
		echo "Error: Proxy not started"
		return 1
	fi

	# Get PID
	PID=$(lsof -i -P | grep ":${LISTEN_PORT} (LISTEN)" | awk -F ' ' '{print $2}')
	if [ "${PID}" == "" ]; then
		echo "Error: Proxy not started"
		return 1
	fi

	# Write TMP file
	if [ ! -d "${TMP_DIR}" ]; then
		mkdir -p "${TMP_DIR}"
	fi
	{
		echo "PID=${PID}"
		echo "PORT=${LISTEN_PORT}"
		echo "PROXY=${PROXY}"
	} > "${TMP_DIR}"/"${LISTEN_PORT}"
}

close() {
	declare ERROR=0
	if [ -d "${TMP_DIR}" ] && [ "$(find "${TMP_DIR}" -maxdepth 0 -empty -exec echo 1 \;)" == "" ]; then
		declare -r FILE=${TMP_DIR}/${1}
		declare -r PORT=$(basename "${FILE}")
		declare -r PID=$(grep 'PID=' "${FILE}" | sed 's/PID=//')
		if ! kill "${PID}"; then
			echo "Error: disconnecting vpn ${PORT}" >&2
			ERROR=1
		fi
		rm "${TMP_DIR}/${PORT}"
	fi
	return "${ERROR}"
}

open_all() {
	declare ERROR=0
	if [ -d "${CONFIG_DIR}/profiles" ]; then
		for FILE in "${CONFIG_DIR}"/profiles/*; do
			PROFILE=$(basename "${FILE}")
			if ! open "${PROFILE}"; then
				ERROR=1
			fi
		done
	fi
	return "${ERROR}"
}

close_all() {
	declare ERROR=0
	if [ -d "${TMP_DIR}" ] && [ "$(find "${TMP_DIR}" -maxdepth 0 -empty -exec echo 1 \;)" == "" ]; then
		for FILE in "${TMP_DIR}"/*; do
			declare LISTEN_PORT=""
			LISTEN_PORT=$(basename "${FILE}")
			if [ "$(close "${LISTEN_PORT}")" == "1" ]; then
				ERROR=1
			fi
		done
	fi
	return "${ERROR}"
}

case ${1} in
	"status") status;;
	"open") open "${2}" "${3}";;
	"close")
		if [ "$(close "${2}")" == "1" ]; then
			exit 1
		fi
	;;
	"open-all")
		if [ "$(open_all)" == "1" ]; then
			exit 1
		fi
	;;
	"close-all")
		if [ "$(close_all)" == "1" ]; then
			exit 1
		fi
	;;
	*)
		showHelp
		exit 1
	;;
esac
