#!/bin/bash
# Apply a command to a subset of a the nodes in a topology

# Initial parameter values
topology=""
src=""
dst=""
nodes="none"
vmusername="root"
portbase=2200
virtnethostpath="`dirname $BASH_SOURCE`/../../"
idfile="`dirname $BASH_SOURCE`/../../data/defaults/root/.ssh/id_rsa"
donotprompt=0

function usage {
	echo "Usage: vn-applycmd CMD [OPTION]..."
	echo "Apply CMD to one or more virtnet nodes."
	echo " "
	echo "Built-in commands are:"
	echo "  cat       cat the file given by --dst on the node(s)"
	echo "  poweroff  shutdown the node(s)"
	echo "  upload    scp file --src on host to --dst on node(s)"
	echo "  test      test the node list, echoing test on the host"
	echo "You may add additional commands (edit this script)."
	echo " "
	echo "Options can be given in long form (--option) or short form (-o). Some options"
	echo "have a parameter, which is always space separated from the option."
	echo "  -n, --nodes NODELIST  list of nodes"
	echo "  -t, --topology NUM    topology number"
	echo "  -s, --src FILE        source file for scp"
	echo "  -d, --dst FILE        dest. file for scp"
	echo "  -y, --yes             do not prompt for confirmation"
	echo "  -h, -?, --help        display this help"
	echo " "
	echo "  FILE is a filename with path"
	echo "  NUM is a topology number"
	echo "  NODELIST is a node number, list of node numbers or one of:"
	echo "    all     all nodes in the topology"
	echo "    host    hosts in the topology (requires nodelist file)"
	echo "    left    left hosts in the topology (requires nodelist file)"
	echo "    right   right hosts in the topology (requires nodelist file)"
	echo "    router  routers in the topology (requires nodelist file)"
	echo "  Using any of the keywords for NODELIST requires --topology NUM option"
	echo "  to be used. Those that require a nodelist file, also require a file"
	echo "  nodelist-topology-NUM to exist and specify node types in the "
	echo "  interfaces directory."
	echo " "
	echo "Examples"
	echo " "
	echo "1. Test that the NODELIST gives the nodes you expect:"
	echo "   vn-applycmd test --nodes \"1 2 3\""
	echo "   This should apply the test on node1, node2 and node3."
	echo " "
	echo "2. Poweroff all nodes in topology 5:"
	echo "   vn-applycmd poweroff --nodes all --topology 5"
	echo " "
	echo "3. Copy a file ./hostname to /etc/hostname on routers in topology 100:"
	echo "   vn-applycmd upload --src ./hostname --dst /etc/hostname --nodes routers --topology 100"
	echo " "
	echo "4. Display contents of file /home/network/out.txt on node 2, without prompting:"
	echo "  vn-applycmd cat --dst /home/network/out.txt --nodes 2 --yes"
	echo " "
 
}

# First option must be the command
if [ "$1" = "" ];
then
	echo "Error. You must issue a command."
	exit 1;
fi
cmd=$1
shift

# Read input parameters
while [ "$1" != "" ]; do
        case $1 in
                -s | --src )   shift;
                                src=$1;
                                ;;
                -d | --dst )  shift;
                                dst=$1;
                                ;;
                -t | --topology )  shift;
				topology=$1;
                                ;;
                -n | --nodes )  shift;
                                nodes=$1;
                                ;;
		-y | --yes )	donotprompt=1;
				;;
                * )             usage
                                exit 1
        esac
        shift
done

# If a topology specified, check that the interface files exist
topologyfilesexist=0
topologyspecified=0
nodelistexists=0
if [ -n "${topology}" ];
then
	topologyspecified=1
	if [ -e ${virtnethostpath}data/topologies/interfaces-topology-${topology}-node-01 ];
	then
		topologyfilesexist=1
	elif [ -e ${virtnethostpath}data/topologies/interfaces-topology-0${topology}-node-01 ];
	then
		topologyfilesexist=1
		topology="0${topology}"
	fi

fi

# Check if nodelist exists
if [ "${topologyfilesexist}" = "1" ];
then
	if [ -e ${virtnethostpath}data/topologies/nodelist-topology-${topology} ];
	then
		nodelistexists=1
	fi
fi


# Get the list of nodes to apply to
if [ "${nodes}" = "none" ];
then
	nodelist="";
elif [ "${nodes}" = "all" -o "${nodes}" = "host" -o "${nodes}" = "router" -o "${nodes}" = "left" -o "${nodes}" = "right" ];
then
	if [ "${topologyfilesexist}" != "1" ];
	then
		echo "Error. If you want to apply to a subset of nodes, such as "
		echo "all, routers or hosts, then you must specify a topology "
		echo "that exists, e.g. --topology 1"
		exit 1;
	else
		if [ "${nodes}" = "all" ];
		then
			nodelist=""
			for i in `ls ${virtnethostpath}data/topologies/interfaces-topology-${topology}-node-*`;
			do
				nodenum=`basename $i | cut -d "-" -f 5`
				if [ `echo ${nodenum} | cut -c 1` = "0" ];
				then
					nodenum=`echo ${nodenum} | cut -c 2`
				fi
				nodelist="${nodelist} node${nodenum}"
			done
		else
			if [ "${nodelistexists}" != "1" ];
			then
				echo "Error. If you want to apply to a subset of nodes, such as "
				echo "all, routers or hosts, then you must ensure the nodelist"
				echo "file exists in the topology directory."
				exit 1;
			else
				if [ "${nodes}" = "host" ];
				then
					pattern="-e left -e right"
				else
					pattern=${nodes}
				fi

				for nodedetails in `grep ${pattern} ${virtnethostpath}data/topologies/nodelist-topology-${topology}`;
				do
					nodelist="${nodelist} node"`echo ${nodedetails} | cut -d "," -f 1`
				done
			fi
		fi
	fi
else
	nodelist=""
	for j in ${nodes}; do
		nodelist="${nodelist} node$j"
	done
fi

sshcmd=0
scpcmd=0
localcmd=0
if [ "${cmd}" = "poweroff" ];
then
	remotecmd="poweroff"
	sshcmd=1
elif [ "${cmd}" = "upload" ];
then
	if [ "${src}" = "" ];
	then
		echo "src must be a file that exists on the host computer."
		exit 1;
	fi
	if [ "${dst}" = "" ];
	then
		echo "dst must be a file on the node"
		exit 1;
	fi
	localfile=${src}
	remotefile=${dst}
	scpcmd=1
elif [ "${cmd}" = "cat" ];
then
	if [ "${dst}" = "" ];
	then
		echo "dst must be a file on the node"
		exit 1;
	fi
	remotecmd="cat ${dst}"
	sshcmd=1

elif [ "${cmd}" = "test" ]; then
	hostcmd="echo test"
	localcmd=1

elif [ "${cmd}" = "YOUR_COMMAND_NAME" ];
then
	# Write your command here
	remotecmd="ls /home/network"
	sshcmd=1
else
	echo "Error. You must use one of the pre-defined commands."
	echo "You can create your own command by editing this file."
	exit 1
fi

if [ "${donotprompt}" = "0" ]; then
	echo "About to apply '${cmd}' to the following nodes:"
	echo ${nodelist}
	echo "Are you ready to proceed? y/n"

	read proceed
else
	proceed="y"
fi

if [ "${proceed}" != "y" ]; then
	echo "Exiting."
	exit 1
fi

for i in ${nodelist};
do
	echo "${i}: ${cmd}"

	nodenum=`echo ${i} | cut -c 5-`
	port=$(( ${portbase} + ${nodenum} ));

	if [ "${scpcmd}" = "1" ];
	then
		scp -o "StrictHostKeyChecking no" -i ${idfile} -P ${port} ${localfile} ${vmusername}@127.0.0.1:${remotefile}
	fi

	if [ "${sshcmd}" = "1" ];
	then
		ssh -o "StrictHostKeyChecking no" -l ${vmusername} -i ${idfile} -p ${port} 127.0.0.1 ${remotecmd}
	fi

	if [ "${localcmd}" = "1" ];
	then
		${hostcmd}
	fi
done