#!/bin/sh

# kvm-simple-init
# Description : Simple init script to manage KVM virtual machines
# URL : https://github.com/tmartinfr/kvm-simple-init
# Copyright 2013 Thomas Martin <thomas@oopss.org>

### BEGIN INIT INFO
# Provides:          kvm-simple-init
# Required-Start:    $remote_fs $syslog
# Required-Stop:     $remote_fs $syslog
# Default-Start:     2 3 4 5
# Default-Stop:      0 1 6
# Short-Description: KVM init script
# Description:       Simple init script to manage KVM virtual machines
### END INIT INFO

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

set -e


###############################################################################
# Hard-coded parameters
###############################################################################

# Configuration directory
CONFDIR=/etc/kvm-simple-init

# Number of seconds before :
WAIT_FOR_START=5        # VM start
WAIT_FOR_SOFT_STOP=60   # VM soft shutdown
WAIT_FOR_HARD_STOP=5    # VM hard shutdown
WAIT_FOR_TEST=2         # VM testing

# Number of seconds before sending command to qemu after connecting to port
QEMU_CMD_WAIT=2

# Timeout before netcat timeout when sending commands to qemu
NC_TIMEOUT=5


###############################################################################
# Help
###############################################################################

usage() {
	cat <<EOT
$0 {start|stop|restart|status} [VM]
EOT
	exit 1
}


###############################################################################
# Functions
###############################################################################

# Test if VM is running or stopped by checking availability of the monitor port.
# Wait until $timeout seconds.
vm_state() {
	state_to_check=$1   # "running" or "stopped"
	monitor_port=$2
	timeout=$3

	while [ $timeout -gt 0 ]; do
		if nc -z localhost $monitor_port; then
			[ "$state_to_check" = "running" ] && return 0
		else
			[ "$state_to_check" = "stopped" ] && return 0
		fi
		sleep 1
		timeout=$(($timeout - 1))
	done

	return 1
}

# Send command to VM using QEMU monitor port
send_qemu_cmd() {
	monitor_port="$1"
	cmd="$2"

	(sleep $QEMU_CMD_WAIT; echo "$cmd";) \
	   | timeout $NC_TIMEOUT nc localhost $monitor_port >/dev/null 2>/dev/null &
}

# Load environment variables in the configuration file
load_vm_config() {
	vmname="$1"
	vmfile="$CONFDIR/$vmname"

	if ! [ -r "$vmfile" ]; then
		echo "$vmfile not readable. Aborting."
		return 1
	fi

	. $vmfile

	[ "$MONITOR_PORT" ] || { echo "$vmfile: Missing MONITOR_PORT."; return 1; }
	[ "$KVM_OPTS" ] || { echo "$vmfile: Missing KVM_OPTS."; return 1; }
}

# Clean global variables imported from the configuration file
cleanup_env() {
	unset MONITOR_PORT KVM_OPTS
}

# Start the VM with the name passed in argument
vm_start() {
	vmname="$1"
	load_vm_config $vmname || return 1

	if vm_state running $MONITOR_PORT $WAIT_FOR_TEST; then
		echo "Not starting $vmname. Already running."
		return 0
	fi

	echo -n "Starting $vmname... "
	kvm $KVM_OPTS

	if vm_state running $MONITOR_PORT $WAIT_FOR_START; then
		echo "Done."
	else
		echo "Failed."
		return 1
	fi
}

# Stop the VM with the name passed in argument
vm_stop() {
	vmname="$1"
	load_vm_config $vmname || return 1

	if vm_state stopped $MONITOR_PORT $WAIT_FOR_TEST; then
		echo "Not stopping $vmname. Already stopped."
		return 0
	fi

	echo -n "Stopping $vmname... "
	send_qemu_cmd $MONITOR_PORT "system_powerdown"

	if vm_state stopped $MONITOR_PORT $WAIT_FOR_SOFT_STOP; then
		echo "Done."
	else
		echo "Failed."

		echo -n "Killing $vmname... "
		send_qemu_cmd $MONITOR_PORT "quit"

		if vm_state stopped $MONITOR_PORT $WAIT_FOR_HARD_STOP; then
			echo "Done."
		else
			echo "Failed."
			return 1
		fi
	fi
}

# Get the status of the VM with the name passed in as an argument
vm_status() {
	vmname="$1"
	load_vm_config $vmname || return 1

	if vm_state running $MONITOR_PORT $WAIT_FOR_TEST; then
		echo "$vmname is running."
	else
		echo "$vmname is stopped."
	fi
}

###############################################################################
# Actions
###############################################################################

if [ $# -eq 1 ]; then
	vmnames=$(ls $CONFDIR)
elif [ $# -eq 2 ]; then
	vmnames=$2
else
	usage
fi

action="$1"
retcode=0

for vmname in $vmnames; do
	if [ "$action" = "start" ]; then
		vm_start $vmname || retcode=1
	elif [ "$action" = "stop" ]; then
		vm_stop $vmname || retcode=1
	elif [ "$action" = "restart" ] || [ "$action" = "force-reload" ]; then
		vm_stop $vmname || retcode=1
		vm_start $vmname || retcode=1
	elif [ "$action" = "status" ]; then
		vm_status $vmname || retcode=1
	else
		usage
	fi

	cleanup_env
done

exit $retcode

