#!/bin/bash
#
# /etc/init.d/virtdomains
# Start / stop guests automatically when host boots / shuts down.
#
# chkconfig: 345 99 00
# description: Start / stop Xen/KVM guests.
#
# This script offers fairly basic functionality.  It should work on Redhat
# but also on LSB-compliant SuSE releases and on Debian with the LSB package
# installed.  (LSB is the Linux Standard Base)
#
# Based on the example in the "Designing High Quality Integrated Linux
# Applications HOWTO" by Avi Alkalay
# <http://www.tldp.org/HOWTO/HighQuality-Apps-HOWTO/>
#
# Revised by T.Nonogaki, Stray Penguin
# <http://www.asahi-net.or.jp/~aa4t-nngk/>
# Rev.0.1.2
#
### BEGIN INIT INFO
# Provides:          virtdomains
# Required-Start:    $syslog $remote_fs libvirtd
# Should-Start:
# Required-Stop:     $syslog $remote_fs libvirtd
# Should-Stop:
# Default-Start:     3 4 5
# Default-Stop:      0 1 2 6
# Default-Enabled:   yes
# Short-Description: Start/stop xen/KVM guests
# Description:       Start / stop guests automatically when host 
#                    boots / shuts down.
### END INIT INFO

# Correct exit code would probably be 5, but it's enough 
# if xend complains if we're not running as privileged domain
#if ! [ -e /proc/xen/privcmd ]; then
#	exit 0
#fi

service libvirtd status > /dev/null || exit 0

LOCKFILE=/var/lock/subsys/virtdomains
XENDOM_CONFIG=/etc/sysconfig/virtdomains

test -r $XENDOM_CONFIG || { echo "$XENDOM_CONFIG not existing";
	if [ "$1" = "stop" ]; then exit 0;
	else exit 6; fi; }

. $XENDOM_CONFIG

# Use the SUSE rc_ init script functions;
# emulate them on LSB, RH and other systems
if test -e /etc/rc.status; then
    # SUSE rc script library
    . /etc/rc.status
else    
    _cmd=$1
    declare -a _SMSG
    if test "${_cmd}" = "status"; then
	_SMSG=(running dead dead unused unknown)
	_RC_UNUSED=3
    else
	_SMSG=(done failed failed missed failed skipped unused failed failed)
	_RC_UNUSED=6
    fi
    if test -e /etc/init.d/functions; then
	# REDHAT
	. /etc/init.d/functions
	echo_rc()
	{
	    #echo -n "  [${_SMSG[${_RC_RV}]}] "
	    if test ${_RC_RV} = 0; then
		success "  [${_SMSG[${_RC_RV}]}] "
		echo
	    else
		failure "  [${_SMSG[${_RC_RV}]}] "
		echo
	    fi
	}
    elif test -e /lib/lsb/init-functions; then
	# LSB    
    	. /lib/lsb/init-functions
        if alias log_success_msg >/dev/null 2>/dev/null; then
	  echo_rc()
	  {
	       echo "  [${_SMSG[${_RC_RV}]}] "
	  }
        else
	  echo_rc()
	  {
	    if test ${_RC_RV} = 0; then
		log_success_msg "  [${_SMSG[${_RC_RV}]}] "
	    else
		log_failure_msg "  [${_SMSG[${_RC_RV}]}] "
	    fi
	  }
        fi
    else    
	# emulate it
	echo_rc()
	{
	    echo "  [${_SMSG[${_RC_RV}]}] "
	}
    fi
    rc_reset() { _RC_RV=0; }
    rc_failed()
    {
	if test -z "$1"; then 
	    _RC_RV=1;
	elif test "$1" != "0"; then 
	    _RC_RV=$1; 
    	fi
	return ${_RC_RV}
    }
    rc_check()
    {
	return rc_failed $?
    }	
    rc_status()
    {
	rc_failed $?
	if test "$1" = "-r"; then _RC_RV=0; shift; fi
	if test "$1" = "-s"; then rc_failed 5; echo_rc; rc_failed 3; shift; fi
	if test "$1" = "-u"; then rc_failed ${_RC_UNUSED}; echo_rc; rc_failed 3; shift; fi
	if test "$1" = "-v"; then echo_rc; shift; fi
	if test "$1" = "-r"; then _RC_RV=0; shift; fi
	return ${_RC_RV}
    }
    rc_exit() { exit ${_RC_RV}; }
    rc_active() 
    {
	if test -z "$RUNLEVEL"; then read RUNLEVEL REST < <(/sbin/runlevel); fi
	if test -e /etc/init.d/S[0-9][0-9]${1}; then return 0; fi
	return 1
    }
fi

if ! which usleep &>/dev/null
then
  usleep()
  {
    if [ -n "$1" ]
    then
      sleep $(( $1 / 1000000 ))
    fi
  }
fi

# Reset status of this service
rc_reset

##
# Returns 0 (success) if the given parameter name is a directory, and
# that directory is not empty.
contains_something()
{
  if [ -d "$1" ] && [ `/bin/ls $1 | wc -l` -gt 0 ]
  then
    return 0
  else
    return 1
  fi
}

# Assume the first non empty, non comment line of the given file
# reads a domain name.
# Returns 0 (success) if a line is found, 1 if the given argument is
# not a file, 2 if the file doesn't contain a valid line.
get_dn_from_f()
{
  local dn=""
  [ -f "$1" ] || return 1
  dn=$(cat $1 | /bin/grep -Ev '^[[:space:]]*(#|$)' | /usr/bin/head -n 1)
  [ -z "$dn" ] && return 2
  echo $dn
  return 0
}

# Read a name from virsh output.
rdname()
{
  local nm
  nm=$(env LANG=C virsh dominfo ${1##*/} 2>/dev/null |
         /bin/awk '/^Name:[[:space:]]+/ { print $2; exit; }')
  echo $nm
}

# Generate a string suitable for "case" matching expression from the
# content of the files under XENDOMAINS_AUTO directory.
rdnames()
{
    local dom fl dom nm
    NAMES=
    if ! contains_something "$XENDOMAINS_AUTO"
    then 
	return
    fi
    for fl in $XENDOMAINS_AUTO/*; do
	dom=$(get_dn_from_f $fl)
	nm=$(rdname $dom)
	if test -z "$NAMES"; then 
	    NAMES=$nm;
	else
	    NAMES="$NAMES|$nm"
	fi
    done
}

# Return 0 (true) if the given domain is up (includes running, paused,
# dying or zombies as well). Return 1 (false) if and only if state is
# "shut off".
is_running()
{
    local nm
    nm=$(rdname $1)
    RC=1
    while read LN; do
	read id name state < <(echo "$LN")
	if [ "$id" = "Id" -o -z "${id##-*-}" ]; then continue; fi
	if [ "$id" = "0" -o "$state" = "shut off" ]; then continue; fi
	case $name in
	     $nm)
		RC=0
		break
		;;
	esac
    done < <(env LANG=C virsh list --all 2>/dev/null)
    return $RC
}

# Starts or restores guest domains. 
start() 
{
    local RETVAL fl dom saved_domains autodoms nm
    if [ -f $LOCKFILE ]; then
	echo -n "virtdomains already running (lockfile exists)"
	rc_failed
	return; 
    fi

    [ -d $(dirname "$LOCKFILE") ] || 
      mkdir -p $(dirname "$LOCKFILE")

    if contains_something "$XENDOMAINS_AUTO"
    then
	# Start all domains with files in XENDOMAINS_AUTO.
	# We expect files for auto starting domains to be in XENDOMAINS_AUTO
	# - each file must contain a line that match the name of the domain.
	# You can manage the starting order through file names, since this
	# sorts the files before processing.
	# NOTE: This order is respected only when STARTing domains.
	# RESTORing doesn't take this into account for now.
	#
	# TODO: We should record which domain name belongs 
	# so we have the option to selectively shut down / migrate later
	# If a domain statefile from $XENDOMAINS_SAVE matches a domain name
	# in $XENDOMAINS_AUTO, do not try to start that domain; if it didn't 
	# restore correctly it requires administrative attention.
	saved_domains=" "
	if [ "$XENDOMAINS_RESTORE" = "true" ] &&
	  contains_something "$XENDOMAINS_SAVE"; then
	    rdnames
	    echo -n "Restoring guest domains: "
	    saved_domains=`ls $XENDOMAINS_SAVE`
	    for dom in $XENDOMAINS_SAVE/*; do
		case ${dom##*/} in
		    $NAMES)
			# do nothing
			;;
		    *)
			continue
			;;
		esac
		echo -n "${dom##*/} "
		virsh restore $dom &>/dev/null
		RETVAL=$?
		if [ $RETVAL -ne 0 ]; then
		    rc_failed $RETVAL
		    echo -n '!'
		else
		    touch $LOCKFILE
		    #mv $dom ${dom%/*}/.${dom##*/}
		    #rm $dom
		fi
		echo -n ". "
	    done
	    HAVE_SAVE=1
	fi

	[ $HAVE_SAVE -eq 1 ] && echo
	echo -n "Starting auto guest domains: "
	autodoms=($(ls -1 $XENDOMAINS_AUTO |sort))
	for fl in ${autodoms[@]}; do
	    nm=$(get_dn_from_f ${XENDOMAINS_AUTO}/$fl)
	    if [ $? -ne 0 ]; then
		continue
	    fi
	    echo -n "$nm"
	    echo $saved_domains | grep -qw $nm >/dev/null
	    if [ $? -eq 0 ]; then
		echo -n "(skip)"
	    elif is_running $nm; then
		echo -n "(skip)"
	    else
		echo -n " "
		virsh start $nm &>/dev/null
		RETVAL=$?
		if [ $RETVAL -ne 0 ]; then
		    rc_failed $RETVAL
		    echo -n '!'
		else
		    touch $LOCKFILE
		    usleep $XENDOMAINS_CREATE_USLEEP
		fi
		HAVE_AUTO=1
	    fi
	    echo -n ". "
	done
	[ $HAVE_AUTO -eq 1 ] || echo
    fi	
}

# Return TRUE if it seems only zombie domains (not in shutdown process)
# left, FALSE if no domains left or found any shutting-down healthily.
# FIXME: Are these check conditions correct?
all_zombies()
{
    local all=0
    local zomb=0
    while read LN; do
	read id name state < <(echo "$LN")
	if [ "$id" = "Id" -o -z "${id##-*-}" ]; then continue; fi
	if [ "$id" = "0" -o "$state" = "shut off" ]; then continue; fi
	: $((all++))
	case $state in
	    dying)
		: $((zomb++))
		;;
	    *)
		# do nothing
		;;
	esac
    done < <(env LANG=C virsh list --all 2>/dev/null)
    [ $all -eq 0 ] && return 1
    [ $all -eq $zomb ] && return 0
    return 1
}

# Call syntax: watchdog_vm $dom [force]
# Wait for max $XENDOMAINS_STOP_MAXWAIT for a guest to shutdown;
# if it has not exited by that time kill it, so the init script will
# succeed within a finite amount of time; if force is added,
# it will kill the guest immediately.
watchdog_vm()
{
    local dom PSAX opt
    local force=0
    : ${XENDOMAINS_STOP_MAXWAIT:=0}
    dom=$1
    if [ -n "$2" ]; then
	[ $2 = force ] && force=1
    fi
    usleep 300000
    for no in `seq 0 $XENDOMAINS_STOP_MAXWAIT`; do
	# go to kill immediately if force flag was set
	if [ $force -eq 1 ]; then
	    break
	fi
	# exit if save/migrate/shutdown is finished
	if ! is_running $dom; then return; fi
	echo -n "."; sleep 1
    done
    sleep 1

    # If the domain still left -
    # give it another polite try
    is_running $dom || return 0
    echo -n "(shut)."
    virsh shutdown $dom &>/dev/null

    # then forcibly destroy
    is_running $dom || return 0
    echo -n "(destroy)."
    virsh destroy $dom &>/dev/null
    is_running $dom || return 0

    # finally kill the qemu-kvm process
    PSAX=$(ps axlww | grep qemu-kvm | grep "-name $dom" | grep -v grep)
    read PSF PSUID PSPID PSPPID < <(echo "$PSAX")
    echo -n "(TERM)."
    kill $PSPID &>/dev/null
    PSAX=$(ps axlww | grep qemu-kvm | grep "-name $dom" | grep -v grep)
    test -z "$PSAX" && return 0
    echo -n "(KILL)."
    kill -s KILL $PSPID &>/dev/null
    return 0
}

stop()
{
    local RETVAL id name state
    # Collect list of domains to shut down
    if test "$XENDOMAINS_AUTO_ONLY" = "true"; then
	rdnames
    fi
    echo -n "Shutting down guest domains: "
    if ! check_domains_left; then
	rm -f $LOCKFILE
	RETVAL=0
	return 0
    fi
    while read LN; do
	read id name state < <(echo "$LN")
	if [ "$id" = "Id" -o -z "${id##-*-}" ]; then continue; fi
	if [ "$id" = "0" -o -z "$id" ]; then continue; fi
	if [ "$state" != "running" ]; then continue; fi
	echo -n "$name"
	if test "$XENDOMAINS_AUTO_ONLY" = "true"; then
	    case $name in
		$NAMES)
		    # nothing
		    ;;
		*)
		    echo "(skip). "
		    continue
		    ;;
	    esac
	fi

	if test "$state" = "dying"; then
	    echo -n "(zomb). "
	    continue
	fi

	if test -n "$XENDOMAINS_MIGRATE"; then
	    echo -n "(migr)"
	    virsh migrate $name $XENDOMAINS_MIGRATE &>/dev/null
	    RETVAL=$?
	    if [ $RETVAL -ne 0 ]; then
		rc_failed $RETVAL
		echo -n '!. '
	    else
		watchdog_vm $name
		echo -n ". "
		continue
	    fi
	fi

	if test -n "$XENDOMAINS_SAVE"; then
	    echo -n "(save)"
	    mkdir -p "$XENDOMAINS_SAVE"
	    virsh save $name ${XENDOMAINS_SAVE}/$name &>/dev/null
	    RETVAL=$?
	    if [ $RETVAL -ne 0 ]; then
		rc_failed $RETVAL
		echo -n '!. '
	    else
		watchdog_vm $name
		echo -n ". "
		continue
	    fi
	fi

	if test -n "$XENDOMAINS_SHUTDOWN"; then
	    # XENDOMAINS_SHUTDOWN should be "--halt --wait"
	    echo -n "(shut)"
	    virsh shutdown $name &>/dev/null
	    RETVAL=$?
	    if [ $RETVAL -ne 0 ]; then
		rc_failed $RETVAL
		echo -n '!'
	    else
		watchdog_vm $name
	    fi
	    echo -n ". "
	fi
    done < <(env LANG=C virsh list --all 2>/dev/null)

    # NB. this stops ALL guest domains, not just the ones in AUTODIR/*
    # This is because it's easier to do ;-) but arguably if this script is run
    # on system shutdown then it's also the right thing to do.
    if test -n "$XENDOMAINS_SHUTDOWN_ALL"; then
	if ! all_zombies; then
	    # XENDOMAINS_SHUTDOWN_ALL should be "--all --halt --wait"
	    if ! check_domains_left; then
		rm -f $LOCKFILE
		return
	    fi
	    echo
	    echo -n "SHUTDOWN_ALL "
	    while read LN; do
		read id name state < <(echo "$LN")
		if [ "$id" = "Id" -o -z "${id##-*-}" ]; then continue; fi
		if [ "$id" = "0" -o -z "$id" ]; then continue; fi
		if [ "$state" = "shut off" ]; then continue; fi
		echo -n "$name"
		if [ "$state" = "running" ]; then
		    echo -n "(shut)"
		    virsh shutdown $name &>/dev/null
		    RETVAL=$?
		    if [ $RETVAL -ne 0 ]; then
			rc_failed $RETVAL
			echo -n '! '
		    else
			watchdog_vm $name
			echo -n ". "
		    fi
		else
		    watchdog_vm $name force
		    echo -n ". "
		fi
	    done < <(env LANG=C virsh list --all 2>/dev/null)
	fi
    fi

    # Unconditionally delete lock file
    rm -f $LOCKFILE
}

check_domain_up()
{
    while read LN; do
	read id name state < <(echo "$LN")
	if [ "$id" = "Id" -o -z "${id##-*-}" ]; then continue; fi
	if [ "$id" = "0" -o -z "$id" ]; then continue; fi
	if [ "$state" != "running" ]; then continue; fi
	case $name in 
	    $1)
		return 0
		;;
	esac
    done < <(env LANG=C virsh list --all 2>/dev/null)
    return 1
}

check_domains_left()
{
    while read LN; do
	read id name state < <(echo "$LN")
	if [ "$id" = "Id" -o -z "${id##-*-}" ]; then continue; fi
	if [ "$id" = "0" -o -z "$id" ]; then continue; fi
	case $state in
	    "shut off")
		continue;;
	    *)
		return 0
		;;
	esac
    done < <(env LANG=C virsh list --all 2>/dev/null)
    return 1
}

check_all_auto_domains_up()
{
    local fl nm dom missing
    if ! contains_something "$XENDOMAINS_AUTO"
    then
      return 0
    fi
    missing=
    for fl in $XENDOMAINS_AUTO/*; do
	dom=$(get_dn_from_f $fl)
	[ $? -ne 0 ] && continue
	nm=$(rdname $dom)
	if check_domain_up "$nm"; then 
	    echo -n " $nm"
	else 
	    missing="$missing $nm"
	fi
    done
    if test -n "$missing"; then
	echo -n " MISS AUTO:$missing"
	return 1
    fi
    return 0
}

check_all_saved_domains_up()
{
    local missing
    if ! contains_something "$XENDOMAINS_SAVE" 
    then
      return 0
    fi
    missing=`/bin/ls $XENDOMAINS_SAVE`
    echo -n " MISS SAVED: " $missing
    return 1
}

# This does NOT necessarily restart all running domains: instead it
# stops all running domains and then boots all the domains specified in
# AUTODIR.  If other domains have been started manually then they will
# not get restarted.
# Commented out to avoid confusion!

restart()
{
    stop
    rc_status -v
    start
    rc_status
}

reload()
{
    restart
}

HAVE_SAVE=0
HAVE_AUTO=0

case "$1" in
    start)
	start
	rc_status
	if test -f $LOCKFILE; then rc_status -v; fi
	;;

    stop)
	stop
	rc_status -v
	;;

    restart)
	restart
	;;
    reload)
	reload
	;;

    status)
	echo -n "Checking for virtdomains:" 
	if test ! -f $LOCKFILE; then 
	    rc_failed 3
	else
	    check_all_auto_domains_up
	    rc_status
	    #check_all_saved_domains_up
	    #rc_status
	fi
	rc_status -v
	;;

    *)
	echo "Usage: $0 {start|stop|restart|reload|status}"
	rc_failed 3
	rc_status -v
	;;
esac

rc_exit