#!/bin/env bash

# Script starts or stops KVM machines or networks in your environment through virsh.
# I.e. your environment must be set up already so you
# - either are able to administer the system's KVM environment (qemu:///system), or
# - have your environment variables set up so you aim at the right KVM environment
#
# Future versions may explicitly ask for the QEMU URI variable and default to qemu:///system, but for
# now we take this for granted that virsh works right for you without specific connection parameters.
#
# The maximum time in seconds to wait for the net-* commands to be successful:
KHNTWAIT=10
# The maximum time in seconds to wait for the machine commands to be successful:
KHVMWAIT=60
# The interval in seconds a network should be polled for status change when executing a start/shutdown -
# while it is called *NT*, this will also be used for starting a VM:
KHNTIVAL="5"
# The interval in seconds a VM should be polled for status change when executing a shutdown:
KHVMIVAL="10"

##### SCRIPT FROM HERE #####


export SHUDDUP=0
prsopt=$(getopt -n "$0" -o q -- "$@")
eval "set -- $prsopt"
while [ "$#" -gt 0 ];do
	case "$1" in
		(-q) SHUDDUP=1;shift;;
		(--) shift;break;;
		(*) exit 1;;
	esac
done

RETVAL=0
KHCMD="$1"
shift
KHTARGS=( "$@" )
declare -x INFOCMD STARTCMD STOPCMD ACTGREP ACTYES ACTNO KILLCMD QWAIT QSTATE
export QLISTNPARMS=( "--name" "--all" )

command -V virsh >/dev/null||exit 255
command -V tput >/dev/null||exit 255

kh_help() {
	[ "$SHUDDUP" -eq 1 ]&&return
	echo -n "USAGE: $(basename "$0") vm-start|vm-stop|net-start|net-stop"
	echo " VMNAME|NETNAME [VMNAME|NETNAME ...]"
}

case "$KHCMD" in
	"net-start"|"net-stop"|"vm-start"|"vm-stop")
		[ "${#KHTARGS[@]}" -eq 0 ]&&echo "Please specify a target."&&kh_help&&exit 127
	;;
	"")
		echo "Please specify a command."
		kh_help
		exit 127
	;;
	*)
		echo "$KHCMD is not a valid command." >&2
		kh_help
		exit 127
	;;
esac

kh_status() {
	# Function should be called with exit status of stop/start command or
	# "start name"/"stop name" for the initial message.
	case "$1" in
		"start")
			[ "$SHUDDUP" -eq 0 ]&&echo -ne "$(tput cub 666)[....] Starting $2...\\033[s\\033[K"||:
		;;
		"stop")
			[ "$SHUDDUP" -eq 0 ]&&echo -ne "$(tput cub 666)[....] Shutting down $2...\\033[s\\033[K"||:
		;;
		255)
			[ "$SHUDDUP" -eq 0 ]&&echo -ne "$(tput cub 666)[$(tput setaf 3)WARN$(tput sgr0)]"
			[ "$SHUDDUP" -eq 0 ]&&echo -e "\\033[u\\033[K successfully force-killed. ($(tput setaf 3)warn$(tput sgr0))"||:
		;;
		0)
			[ "$SHUDDUP" -eq 0 ]&&echo -ne "$(tput cub 666)[$(tput setaf 2) OK $(tput sgr0)]"
			[ "$SHUDDUP" -eq 0 ]&&echo -e "\\033[u\\033[K done."||:
		;;
		10)
			[ "$SHUDDUP" -eq 0 ]&&echo -ne "$(tput cub 666)[$(tput setaf 7)DONE$(tput sgr0)]"
			[ "$SHUDDUP" -eq 0 ]&&echo -e "\\033[u\\033[K already done."||:
		;;
		*)
			[ "$SHUDDUP" -eq 0 ]&&echo -ne "$(tput cub 666)[$(tput setaf 1)FAIL$(tput sgr0)]"
			[ "$SHUDDUP" -eq 0 ]&&echo -e "\\033[u\\033[K $(tput setaf 1)failed$(tput sgr0)."||:
		;;
	esac
}
kh_qstate() {
	case "$ACTGREP" in
		"") QSTATE="$(virsh "$INFOCMD" "$2" 2>/dev/null)";;
		*) QSTATE="$(virsh "$INFOCMD" "$2" 2>&1|grep "$ACTGREP"|awk '{print $NF}')";;
	esac
	case "$1" in
		"start")
			[ "$QSTATE" == "$ACTYES" ]&&return 0||return 1
		;;
		"stop")
			[ "$QSTATE" == "$ACTNO" ]&&return 0||return 1
		;;
	esac
}
kh_exec() {
	# Function expects exactly two arguments: command and KVM net name to be started/stopped.
	# RC:
	# 	1 if no argument passed or arguments invalid
	# 	2 if start is unsuccessful or times out.
	FN="kh_exec"
	[ -z "$1" ]&&echo "Critical exception in $FN(): no argument passed!" >&2&&exit 1
	[ -z "$2" ]&&echo "Critical exception in $FN(): no network name passed!" >&2&&exit 1
	if ! virsh "$LISTCMD" "${QLISTNPARMS[@]}"|grep "$2" >/dev/null;then
		echo "Critical exception in $FN(): Network/VM $2 unknown!" >&2&&exit 2
	fi
	case "$1" in
		"start") kh_status start "$2";;
		"stop") kh_status stop "$2";;
	esac
	case "$1" in
		"start")
			kh_qstate start "$2"&&kh_status 10 "$2"&&return 0
		;;
		"stop")
			kh_qstate stop "$2"&&kh_status 10 "$2"&&return 0
		;;
	esac
	TOELAPSE="0"
	while [ "$TOELAPSE" -lt "$QWAIT" ];do
		case "$1" in
			"start") virsh "$STARTCMD" "$2" >/dev/null 2>&1;;
			"stop") virsh "$STOPCMD" "$2" >/dev/null 2>&1;;
		esac
		TOELAPSE=$((TOELAPSE+QIVAL))
		sleep "$QIVAL"
		[ "$SHUDDUP" -eq 0 ]&&echo -n "."
		case "$1" in
			"start")kh_qstate start "$2"&&break;;
			"stop")kh_qstate stop "$2"&&break;;
		esac
	done
	case "$1" in
		"start")
			if kh_qstate start "$2";then
				kh_status 0
			else
				kh_status 1;RETVAL=20
			fi
		;;
		"stop")
			if kh_qstate stop "$2";then
				kh_status 0
			else
				virsh "$KILLCMD" "$2" >/dev/null 2>&1
				if kh_qstate stop "$2";then
					kh_status 255;RETVAL=255
				else
					kh_status 1;RETVAL=20
				fi
			fi
		;;
	esac
}
kh_filter() {
	# Function expects exactly two arguments: command and KVM net name to be started/stopped.
	# This works as a pre-filter for knstart and knstop
	# RC:
	# 	1 if no argument, no network name passed, or wrong syntax
	# 	2 if network/VM does not exist
	#	3 if network/VM cannot be started
	#	4 if network/VM state is not determinable
	FN="kh_filter"
	[ -z "$1" ]&&echo "Critical exception in $FN(): no argument passed!" >&2&&exit 1
	[ -z "$2" ]&&echo "Critical exception in $FN(): no action specified!" >&2&&exit 1
	[ -z "$3" ]&&echo "Critical exception in $FN(): no network/VM name passed!" >&2&&exit 1
	case "$1" in
		"net")
			INFOCMD="net-info";	STARTCMD="net-start"
			STOPCMD="net-destroy";	ACTGREP="^Active:"
			ACTYES="yes";		ACTNO="no"
			KILLCMD="$STOPCMD";	LISTCMD="net-list"
			QWAIT="$KHNTWAIT";	QIVAL="$KHNTIVAL"
		;;
		"vm")
			INFOCMD="domstate";	STARTCMD="start"
			STOPCMD="shutdown";	#ACTGREP="^State:"
			ACTYES="running";	ACTNO="shut off"
			KILLCMD="destroy";	LISTCMD="list"
			QWAIT="$KHVMWAIT"	QIVAL="$KHVMIVAL"
			[ "$2" == "start" ]&&QIVAL="$KHNTIVAL"
		;;
	esac
	case "$2" in
		"start") kh_exec start "$3";;
		"stop") kh_exec stop "$3";;
		*) echo "Invalid action $2 to $FN()!" >&2&&exit 1;;
	esac
}

[ ${#KHTARGS} -eq 0 ]&&echo "No VM/network specified, aborting." >&2&&exit 127
for i in "${KHTARGS[@]}";do
	case "$KHCMD" in
		"net-stop") kh_filter net stop "$i"||RETVAL="$((RETVAL+1))";;
		"net-start") kh_filter net start "$i"||RETVAL="$((RETVAL+1))";;
		"vm-stop") kh_filter vm stop "$i"||RETVAL="$((RETVAL+1))";;
		"vm-start") kh_filter vm start "$i"||RETVAL="$((RETVAL+1))";;
	esac

done

exit "$RETVAL"