#!/bin/sh

SUDO=0
SUDO_BIN=sudo
EXEC=""
SSHONCE=0
TRUNCATE=0
TEMPDIR=""

usage() {
	echo "$0 [-p] [-d] [-s [-e sudo|doas]] server [...]"
	exit 0
}

# $1 = directory name
# $2 = remote server
# $3 = tempdir
copy_files() {
	# -l = keep symlink / -D = special device
	if [ -d "${1}" ]
	then
		LIST=$(mktemp /tmp/drist-rsync.XXXXXXXXXX)
		if [ -f "$LIST" ]
		then
			printf 'Copying files from "%s" to temporary directory %s:\n' "$1" "$3"
			find "${1}"/ -type f -or -type l | cut -d '/' -f 2- | tee "${LIST}" | sed 's/^/	\//'
			rsync -t -e "ssh $SSH_PARAMS" -lD --files-from="${LIST}" "${1}/" "${2}":"/${3}"
			rm "$LIST"
		fi
	fi
}

# $1 = script filename
# $2 = remote server
# $3 = tempdir
remote_script() {
	if [ -f "${1}" ]
	then
		printf 'Executing file "%s":\n' "$1"
		ssh $SSH_PARAMS "${2}" "cd ${3} && DRIST=${3}/script &&
		    cat - > \$DRIST &&
		    chmod u+x \$DRIST &&
		    ${EXEC} \$DRIST" < "$1"
    if [ "$?" -ne 0 ]; then
      printf '\e[31m\n\n========== EXECUTION OF SCRIPT FAILED for server %s ==========\n\n\e[0m' "${2}"
      exit 1;
    fi
	fi
}

# $1 = remote server
create_temp() {
    TEMPDIR=$(ssh $SSH_PARAMS "$1" "mktemp -d ~/.drist_files_XXXXXXXXXXXXXXX")
    if [ "$TEMPDIR" = "" ]; then
        echo "mktemp error, aborting"
        exit 1
    fi
}

# $1 = remote server
# $2 = temporary directory
delete_temp() {
    if echo "${2}" | grep drist_files_ >/dev/null ; then
        ssh $SSH_PARAMS "$1" "rm -fr ${2}"
    else
        echo "Problem, TEMPDIR was reset during execution, current value is = $2"
	exit 2
    fi
}


# RUNTIME BEGINS HERE
while getopts pndse: arg; do
	case ${arg} in
		d) TRUNCATE=1 ;;
		s) SUDO=1 ;;
		e) SUDO_BIN="${OPTARG}" ;;
		p) SSHONCE=1 ;;
		*) usage ;;
	esac
done
shift $((OPTIND - 1))

# allow to use a privilege escalation program
if [ "$SUDO" -eq 1 ]
then
	EXEC="$SUDO_BIN"
fi

# use ControlMaster to make connections persistent
if [ "$SSHONCE" -eq 1 ]
then
	SSH_PARAMS=-o"ControlMaster=auto"" "-o"ControlPath=/tmp/drist_ssh_%h_%p_%r.sock"" "-o"ControlPersist=1m"
fi

# start looping over server list
if [ -f "$1" ]
then
	SERVER_LIST="$(tr '\n' ' ' < $1)"
else
	SERVER_LIST="$@"
fi

if [ "${SERVER_LIST}" = "" ]
then
	echo "No server specified"
	exit 1
fi

for remote_server in ${SERVER_LIST}
do
	echo "Running on ${remote_server}"

	# check if host exists
	HOSTNAME=$(ssh $SSH_PARAMS "${remote_server}" "uname -n")
	if [ "$?" -ne 0 ]; then
		echo "Error while ssh ${remote_server}"
		exit 2
	fi

	if [ "$TRUNCATE" -eq 1 ]; then
		HOSTNAME="${HOSTNAME%%.*}"
	fi

	create_temp "${remote_server}"
	copy_files "files" "${remote_server}" "$TEMPDIR"
	copy_files "files-${HOSTNAME}" "${remote_server}" "$TEMPDIR"
	remote_script "script" "${remote_server}" "$TEMPDIR"
	remote_script "script-${HOSTNAME}" "${remote_server}" "$TEMPDIR"
	delete_temp "${remote_server}" "$TEMPDIR"

	# close socket if persistance is actived
	if [ "$SSHONCE" -eq 1 ]
	then
		ssh $SSH_PARAMS -O exit -N "$1"
	fi

	unset TEMPDIR HOSTNAME
done