#! /bin/bash

if [ "$1" = '-ssh' ]; then
	shift

	RDP_SSH_JUMPBOX="$1"
	shift
else
	RDP_SSH_JUMPBOX=''
fi

RDP_HOST="$1"
shift

function printHelp() {
	echo 'Usage: rdp [-ssh host] <host> [<xfreeRdpArgs...>]'
}

if [ -z "${RDP_HOST}" ]; then
	printHelp >&2

	exit 1
fi

# Include a config file, which is actually just a script
# This script can do things like update the "RDP_SSH_JUMPBOX"
# based on the "RDP_HOST", provide a password (from a password
# manager, like hunter2, ideally), etc.
if [ -f ~/.rdp.conf ]; then
	. ~/.rdp.conf
fi

## Split the target into host/port pairs
case "${RDP_HOST}" in
	*:*)
		hostName="$(echo "${RDP_HOST}" | cut -f 1 -d ':')"
		hostPort="$(echo "${RDP_HOST}" | cut -f 2 -d ':')"
		;;
	*)
		hostName="${RDP_HOST}"
		hostPort='3389'
		;;
esac

# If we are jumping through an SSH jumpbox, setup port forwarding
if [ -n "${RDP_SSH_JUMPBOX}" ]; then
	# Pick a random port to forward on
	randomPort=$[${RANDOM} % 1024 + 3390]

	# Determine an SSH control socket
	sshControlSocket="/tmp/ssh-sock-${randomPort}-$$${RANDOM}${RANDOM}${RANDOM}"

	function cleanup() {
		# Kill the SSH session
		sshPid="$(timeout 30 ssh -S "${sshControlSocket}" -O check /dev/null 2>&1 | grep 'pid=' | sed 's@.*pid=@@;s@).*$@@')"
		if [ -n "${sshPid}" ]; then
			kill -9 "${sshPid}" >/dev/null 2>/dev/null
		fi

		rm -f "${sshControlSocket}"
	}

	# Start SSH
	ssh -L${randomPort}:${hostName}:${hostPort} -o ExitOnForwardFailure=yes -o ControlMaster=yes -S "${sshControlSocket}" -a -f -N "${RDP_SSH_JUMPBOX}"

	# Verify SSH started
	sshPid="$(ssh -S "${sshControlSocket}" -O check /dev/null 2>&1 | grep 'pid=' | sed 's@.*pid=@@;s@).*$@@')"
	if [ -z "${sshPid}" ]; then
		echo "Unable to start SSH" >&2

		exit 1
	fi

	hostId="${hostName}:${hostPort}!${RDP_SSH_JUMPBOX}"

	hostName='localhost'
	hostPort="${randomPort}"
else
	function cleanup() {
		return 0
	}

	hostId="${hostName}:${hostPort}"
fi

trap cleanup EXIT

# Connect to the RDP session
## Determine username and domain
if [ -n "${RDP_USERNAME}" ]; then
	addArgs=("${addArgs[@]}" -u "${RDP_USERNAME}")
else
	addArgs=("${addArgs[@]}" -u "$(whoami)")
fi

if [ -n "${RDP_DOMAIN}" ]; then
	addArgs=("${addArgs[@]}" -d "${RDP_DOMAIN}")
fi

## Determine port
if [ -n "${hostPort}" -a "${hostPort}" != '3389' ]; then
	addArgs=("${addArgs[@]}" -t "${hostPort}")
fi

## Export the password to the environment and tell xfreerdp to use
## it.  This requires a patched xfreerdp.
if [ -n "${RDP_PASSWORD}" ]; then
	export RDP_PASSWORD

	addArgs=("${addArgs[@]}" -p ENV:RDP_PASSWORD)
fi

## Actually call xfreerdp
xfreerdp --plugin drdynvc --plugin rdpdr --data scard "disk:pwd:$(pwd)" -- --certificate-name "${hostId}" "${addArgs[@]}" "$@" "${hostName}"
returnCode="$?"

exit "${returnCode}"