#!/bin/bash

#
# Copyright (C) 2007,2008 Canonical, Ltd.
# Author: Jamie Strandboge <jamie@canonical.com>
# License: GPLv3
#
set -e

ustconf="$HOME/.uqt-vm-tools.conf"
if [ -s "$ustconf" ]; then
    . "$ustconf"
else
    echo "Could not find '$ustconf'"
    exit 1
fi

. $UQT_VM_TOOLS/libvm.sh
abort_if_root

help() {
    echo "Usage: vm-cmd [-f] [-r] [-s] [-t TIMEOUT] -p <prefix> <command>"
    echo "Usage: vm-cmd [-f] [-r] [-s] [-t TIMEOUT] <vm> <command>"
    echo " -p PREFIX      use PREFIX to construct hostname list"
    echo " -s             start the VM (and shutdown if it wasn't running)"
    echo " -t TIMEOUT     how long to wait for VM to come up if -s used"
    echo " -T TIMEOUT     how long to wait for VM to run a command"
    echo " -f             force the SSH keys to be taken"
    echo " -c             continue if command fails (no effect with -T)"
    echo " -r             login to the VM as root"
    echo ""
    echo "Eg:"
    echo "$ vm-cmd -p sec 'sudo reboot'"
    echo ""
    echo "This will reboot (via ssh) all sec-<release>-<arch> machines"
    echo "that are running, where <release> is in '$vm_release_list' and"
    echo "<arch> is in $vm_archs."
    echo ""
    echo "$ vm-cmd clean-jaunty-i386 'sudo apt-get update'"
    echo ""
    echo "This will run 'apt-get update' on the single VM named 'clean-jaunty-i386'"
}

prefix=
start=
force=
force_continue=0
asroot=
timeout=90
cmd_timeout=0
while getopts "chfrst:T:p:" opt
do
    case "$opt" in
        p) prefix="$OPTARG";;
        s) start=1;;
        f) force=1;;
        c) force_continue=1;;
        r) asroot=1;;
        t) timeout="$OPTARG";;
        T) cmd_timeout="$OPTARG";;
        ?|h) help
           exit 1
           ;;
    esac
done
shift $(( $OPTIND - 1 ))

machines=
if [ -n "$prefix" ]; then
    for a in $vm_archs ; do
        for m in `vm_get_machines_by_prefix "${prefix}" "$a"` ; do
            machines="$machines $m"
        done
    done
else
    if [ -z "$1" ]; then
        help
        exit 1
    fi

    machines="$1"
    shift || true
fi

failed=
for m in $machines ; do
        started=""
        echo "----- $m -----"
        if [ -n "$start" ]; then
            if ! vm_running "$m" ; then
                echo "Starting $m ..."
                virsh start "$m" >/dev/null 2>&1
                started=1
            fi
            host=`vm_wait "$m" $timeout`
        else
            host=`vm_ping ${m}`
        fi
        if [ -z "$host" ]; then
            echo "Skipped $m"
            echo ""
            continue
        fi

        # '-t' forces tty allocation (for sudo passwords)
        args="-t -o BatchMode=yes"
        if [ -n "$force" ]; then
            args="-o StrictHostKeyChecking=no $args"
        fi

        if [ -n "$asroot" ]; then
            args="$args -l root"
        fi

        echo "Command: ssh ${args} $host \"$@\""
        if [ $cmd_timeout -eq 0 ]; then
            # Foreground
            ssh $args "$host" "$@" || {
                rc="$?"
                if [ $force_continue -eq 0 ]; then
                    exit $rc
                else
                    echo "FAIL: command exiting with '$rc'" >&2
                    failed="yes"
                fi
            }
        else
            # Background, waiting for it to finish...
            ssh $args "$host" "$@" &
            pid=$!
            seconds=0
            while [ ! -z $(jobs -p) ]; do
                sleep 1
                seconds=$(( $seconds + 1 ))
                if [ $seconds -gt $cmd_timeout ]; then
                    echo "Command timed out, killing $pid ..."
                    kill $pid
                    break
                fi
                # "jobs -p" doesn't seem to update, so trigger a check with -l
                jobs -l >/dev/null
            done
            wait $pid
        fi
        echo ""

        if [ -n "$started" ]; then
            echo "Shutting down $host ..."
            ssh $args -l root "$host" "shutdown -h now" || true
            # wait for the system to actually shut down
            count=0
            while [ "$count" -lt "$timeout" ]; do
                if [ -z $(vm_ping ${m}) ]; then
                    break
                fi
                sleep 1
                count=$(( $count + 1 ))
            done
        fi
done

if [ -n "$failed" ]; then
    echo "FAIL: command(s) exited with non-zero"
    exit 1
fi
