#!/bin/bash

# Public domain
# (c) Etersoft 2015, 2017

docmd()
{
	echo "\$ $@"
	"$@"
}

fatal()
{
	echo -e "Error: $@"
	exit 1
}

# CHECKME: the same like estrlist has ?
# Note: used egrep! write '[0-9]+(first|two)', not '[0-9]\+...'
rhas()
{
	echo "$1" | egrep -q -- "$2"
}

fastssh()
{
	ssh -o "ControlMaster auto" -o "ControlPath ~/tmp/ssh_mux_%h_%p_%r" -o "ControlPersist 4h" "$@"
}

rooter()
{
	fastssh -x rooter@server -p32 "$@"
}

# USER HOST VEID
addkey()
{
	# TODO: use it not via direct path
	docmd /srv/lav/Projects/git-eter/etersoft-admin-essential/rooter/add_eterkey.sh "$@"
}

# USER HOST VEID GROUP
addgroup()
{
	# TODO: use it not via direct path
	docmd /srv/lav/Projects/git-eter/etersoft-admin-essential/rooter/add_group.sh "$@"
}

# USER HOST VEID [KEYUSER]
setkey()
{
	# TODO: use it not via direct path
	docmd /srv/lav/Projects/git-eter/etersoft-admin-essential/rooter/set_eterkey.sh "$@"
}

# TODO: more complexy grep
vzgrep()
{
	[ -z "$1" ] && cat && return
	grep --color -- "$@"
}

normalize_veid()
{
	#local VEID="$1"
	#VEID=$(($VEID))
	#printf "%03d" "$VEID"
	echo "$(($VEID))"
}

is_number()
{
	# FIXME: some hack
	# TODO: tested is_number in epm functions
	echo "$1" | grep -q -- "^[0-9].*" || return
	[ "$(normalize_veid "$1")" != "0" ] || [ "$1" = "0" ]
}

is_alias()
{
	local HOST="$1"
	grep -q -- "^[Hh]ost[[:space:]]*$HOST" ~/.ssh/config && return
}

# copied from rooter
show_alias()
{	local ALIAS="$1"
	grep -A20 -- "[Hh]ost[[:space:]]*$ALIAS\$" ~/.ssh/config | awk 'BEGIN{desk=0}{if(/^[[:space:]]$/){desk++};if(desk==0){print}}'
}

# add host alias to ~/.ssh/config if needed
check_host_in_config()
{
	is_alias "$HOST" && return
	echo "No available alias for $HOST in ~/.ssh/config. Adding..."
	echo >>~/.ssh/config
	rooter hostconfig "$HOST" >> ~/.ssh/config
	echo "Done."
}

get_ip()
{
	IP=
	[ -n "$VEID" ] && IP=$(rooter vzlist --IP "$HOST" "$VEID")
	if [ "$IP" = "-" ] ; then
		IP=""
	fi
	echo "$IP"
}

get_port()
{
	rooter iptables "$HOST" | grep "to:$IP:[23]2$" | sed -e "s|.*tcp dpt:\([0-9]*\) .*|\1|g" | head -n1
}

get_host()
{
	[ -n "$VEID" ] || return
	rooter vzlist --HOSTNAME "$HOST" "$VEID" 2>/dev/null
}

HOSTLIST="$(rooter hostlist)"

host_in_list()
{
	for host in $HOSTLIST ; do
		[ "$host" = "$1" ] && return
	done
	return 1
}


connect_ssh()
{
VEID=$(normalize_veid "$VEID")
IP=$(get_ip)
PORT=$(get_port)
DIRECTHOST=$(get_host)
# FIXME: default port 32?
DIRECTPORT=32
ARGUSER=''
[ -n "$USERNAME" ] && ARGUSER="$USERNAME@"

if [ "$1" = "ubc" ] || [ "$1" = "vzubc" ] ; then
	rooter vzubc $HOST $VEID
	exit
fi

if [ "$1" = "alias" ] ; then
	ALIAS=$(echo "$DIRECTHOST" | sed -e "s|\..*||g")
	if is_alias "$ALIAS" ; then
		echo "Alias $ALIAS is already exists:"
		show_alias "$ALIAS"
		exit 1
	fi

	if [ -z "$PORT" ] ; then
		[ -n "$DIRECTHOST" ] || fatal "Can't get connect port for $HOST:$VEID"
		PORT=$DIRECTPORT
	fi

	rooter hostconfig $HOST | \
		sed -e "s|^\(Host.*\)\($HOST\)|\1$ALIAS|gi" | \
		sed -e "s|\(Port[[:space:]]*\).*|\1$PORT|gi" >$TMPDIR/sshvz.tmp
	if [ -n "$USERNAME" ] ; then
		if grep -qi "User" $TMPDIR/sshvz.tmp ; then
			subst "s|\(User[[:space:]]*\).*|\1$USERNAME|gi" $TMPDIR/sshvz.tmp
		else
			grep -i "Hostname" $TMPDIR/sshvz.tmp | sed -e "s|\(Hostname\)\([[:space:]]*\).*|User\2$USERNAME|gi" >>$TMPDIR/sshvz.tmp
		fi
	fi

	if [ -s $TMPDIR/sshvz.tmp ] ; then
		echo >>~/.ssh/config
		cat $TMPDIR/sshvz.tmp | grep -v "^$" >> ~/.ssh/config
	else
		fatal
	fi

	echo "Created alias in ~/.ssh/config:"
	show_alias "$ALIAS"
	exit
fi

if [ -z "$PORT" ] ; then
	if [ -n "$DIRECTHOST" ] ; then
		echo "Direct connecting to $DIRECTHOST on $HOST... (possible you need add -pPORT)"
		docmd ssh -t -p$DIRECTPORT "$ARGUSER$DIRECTHOST" "$@"
		exit
	fi
	fatal "Can't get connect port for $HOST:$VEID"
fi

check_host_in_config "$HOST"
echo "Connecting via $HOST:$PORT to $DIRECTHOST ($IP) ..."
docmd ssh -t -p "$PORT" "$ARGUSER$HOST" "$@"
}


Help="sshvz - connect via ssh to host machine or container inside it. (c) Etersoft 2015, 2017
Usage: sshvz [command] [HOST] [command] [VEID | CONTAINERWORD]
Examples:
 $ $0               - print hosts list
 $ $0 HOST          - print containers list
 $ $0 [USER@]WORD(s)       - search in all vz containers' info and enter via ssh
 $ $0 [USER@]HOST WORD(s)  - search in host containers's info and enter via ssh
 $ $0 [USER@]HOST VEID - connect to the container via ssh (your username by default)
 $ $0 [USER@]HOST VEID alias - make alias for this container
 $ $0 HOST VEID ubc - print vzubc info for the container
 $ $0 HOST iptables [WORD(s)] - list iptables or print container related strings from it
Admin only:
 $ $0 HOST [VEID] add USER [keyuser] - add access (admins only) for USER [with key from keyuser]
 $ $0 HOST [VEID] admadd USER [keyuser] - add adm (root) access (admins only) for USER [with key from keyuser]
 $ $0 HOST [VEID] key USER [keyuser] - set key for USER [from keyuser], your by default
where VEID - container ID or 0 for host system
"

HOST="$1"
VEID="$2"

USERNAME=''
if echo "$HOST" | grep -q "@" ; then
    # bashism
    IFS="@" read USERNAME HOST <<< "$HOST"
fi

# skip HOST
shift

if [ "$HOST" = "-h" ] || [ "$HOST" = "--help" ]; then
	echo "$Help"
	exit
fi

# TODO: add -a for all container

if [ -z "$HOST" ]; then
	echo "Available hosts:"
	echo -n -e "\t"
	rooter hostlist
	echo "run $0 -h for help"
	exit
fi

# try to connect by container name by default
if ! host_in_list "$HOST" ; then
	CONT="$HOST"
	echo "Search in... "
	tmpfile=$(mktemp)
	for host in $HOSTLIST ; do
		echo -n " $host" >&2
		rooter vzlist -a --PRINTHOST $host 2>/dev/null
	done | vzgrep "$CONT" > $tmpfile

	echo
	cat $tmpfile

	LINES=$(cat $tmpfile | wc -l)
	if [ "$LINES" = 1 ] ; then
		echo "One container found."
		HOST=$(cat $tmpfile | sed -e "s|.* \([a-z0-9.]*\)$|\1|g")
		VEID=$(cat $tmpfile | sed -e "s|^ *\([0-9]*\) .*|\1|g")
		grep -q " running " $tmpfile || fatal "Container $VEID does not running."
		rm -f $tmpfile
		connect_ssh "$@"
		exit
	fi
	rm -f $tmpfile
	if [ "$LINES" = 0 ] ; then
		fatal "Can't find container $CONT on all hosts"
	fi
	echo "rerun with number of the container"
	exit
fi

if [ "$HOST" = "search" ] || [ "$HOST" = "s" ] ; then
	echo "Search in $HOSTLIST..."
	for host in $HOSTLIST ; do
		rooter vzlist -a --PRINTHOST $host 2>/dev/null
	done | vzgrep "$@"
	exit
fi

# skip VEID
shift

if [ "$VEID" = "search" ] || [ "$VEID" = "s" ] ; then
	rooter vzlist -a $HOST | vzgrep "$@"
	exit
fi

if [ "$VEID" = "iptables" ] ; then
	VEID="$1"
	IP=$(get_ip)
	if [ -z "$IP" ] ; then
		rooter iptables "$HOST" | vzgrep "$@"
	else
		rooter iptables "$HOST" | vzgrep "$IP"
	fi
	exit
fi

if [ "$VEID" = "add" ] ; then
	addkey "$1" "$HOST" 0 "$2"
	exit
fi

if [ "$VEID" = "admadd" ] ; then
	addkey "$1" "$HOST" 0 ""
	addgroup "$1" "$HOST" 0 "wheel"
	exit
fi

if [ "$VEID" = "key" ] ; then
	setkey "$1" "$HOST" 0 "$2"
	exit
fi


if [ "$VEID" != "0" ] && ! is_number "$VEID" ; then
	CONT="$VEID"
	tmpfile=$(mktemp)
	rooter vzlist -a $HOST | vzgrep "$CONT" | tee $tmpfile | vzgrep "$CONT"
	LINES=$(cat $tmpfile | wc -l)
	if [ "$LINES" = 1 ] ; then
		echo "One container found."
		VEID=$(cat $tmpfile | sed -e "s|^ *\([0-9]*\) .*|\1|g")
		grep -q " running " $tmpfile || fatal "Container $VEID does not running."
	elif [ "$LINES" = 0 ] ; then
		fatal "Can't find container $VEID on $HOST"
	else
		echo "rerun with number of the container or 0 for entry to host system"
		exit
	fi
	rm -f $tmpfile
fi


# Just list containers
if [ -z "$VEID" ] ; then
	rooter vzlist -a "$HOST"
	echo "(use 0 for enter to the host system)"
	exit
fi

check_host_in_config "$HOST"

# 0: enter in host machine
if [ "$VEID" = "0" ] ; then
	CMD="$@"
	[ -z "$CMD" ] && CMD="sudo su -"
	docmd ssh -t $HOST $CMD
	exit
fi

if [ "$1" = "add" ] ; then
	addkey "$2" "$HOST" "$VEID" "$3"
	exit
fi

if [ "$1" = "admadd" ] ; then
	addkey "$2" "$HOST" "$VEID" ""
	addgroup "$2" "$HOST" "$VEID" "wheel"
	exit
fi

if [ "$1" = "key" ] ; then
	setkey "$2" "$HOST" "$VEID" "$3"
	exit
fi

connect_ssh "$@"

