#!/bin/bash

show_help() {
  c="  `basename $0`"
  echo "Usage:"
  echo "$c help                         # this help"
  echo "$c restart name                 # restart"
  echo "$c shutdown|stop name           # shutdown and sync"
  echo "$c start name                   # sync and start"
  echo "$c rename name newname          # rename guest"
  echo "$c migrate name host [--live]   # live migrate to host"
  echo "$c move name host [newname]     # offline move guest"
  echo "$c list|l [--all]               # list guests"
  echo "$c info|i [--all|name]          # get information about guest"
  echo "$c edit name [-r RAM|-c CPU|-b device] # edit guest"
  echo "$c set [mem|cpu|boot] guest PARAM # set guest parameters"
  echo "$c console|con|c name           # console for guest"
  echo "$c diff name1 name2             # diff virtual machine configs"
  exit 0
}

shutdown_and_restart() {
  set -x
  virsh shutdown $1
  virsh console $1

  sleep 1s
  sync
  echo 3 > /proc/sys/vm/drop_caches
  sync

  if [ "$2" = "restart" ]; then
    sleep 1s
    virsh start $1 --console
    exit 0
  fi

  sleep 3s
  sync
  echo 3 > /proc/sys/vm/drop_caches
}

start_guest() {
  set -x
  sync
  echo 3 > /proc/sys/vm/drop_caches
  sync
  sleep 1s
  
  virsh start $1 --console
}

rename_guest() {
  TMP=`mktemp`
  virsh dumpxml $1 > $TMP
  sed -i "s|<name>$1</name>|<name>$2</name>|" $TMP
  virsh undefine $1
  virsh define $TMP
  rm -f $TMP
}

move_guest() {
  virsh shutdown $1
  virsh console $1
  echo 3 > /proc/sys/vm/drop_caches
  sync
  if [ "$3" ]; then
    M="$3"
  else
    M="$1"
  fi
  ssh root@$2 virsh start $M --console
}

virsh_list() {
  DOM_IDS=`virsh list $@ | awk '/^ *[0-9-]/ { print $2 }'`
  for domain in $DOM_IDS; do
    virsh dumpxml $domain
  done | awk -F"[<>]" '
    BEGIN {
      sum_memory = 0
      sum_current = 0
      printf("%-20s  %16s  %16s  %7s  %9s\n",
             "Name", "Memory", "Current Memory", "Running", "Autostart")
    }
    /^<domain / { running = "no" }
    /^<domain .* id=/ { running = "YES" }
    /^ *<name>/ {
      name = $3
      fn = "/etc/libvirt/qemu/autostart/" name ".xml"
      astart = ""
      getline astart <fn
      if (astart!="") {
        autostart = "YES"
      } else {
        autostart = "no"
      }
    }
    /^ *<memory>/ { memory = $3; sum_memory += $3 }
    /^ *<currentMemory>/ {
      current = $3
      sum_current += $3
      printf("%-20s %14.2f MB %14.2f MB  %7s  %9s\n",
             name, memory/1024, current/1024, running, autostart)
    }
    END {
      printf("----------------------------------------------------------------------------\n")
      printf("%-20s %14.2f MB %14.2f MB\n",
             "TOTAL:", sum_memory/1024, sum_current/1024)
    }'
}

virsh_info() {
  if [ "$2" == "--all" ]; then
    DOM_IDS=`virsh list --all | awk '/^ *[0-9-]/ { print $2 }'`
  elif [ "$2" ]; then
    DOM_IDS="$2"
  else
    DOM_IDS=`virsh list | awk '/^ *[0-9-]/ { print $2 }'`
  fi
  for domain in $DOM_IDS; do
    virsh dumpxml $domain
  done | awk -F"[<>]" '
    function get0(i) {
      split($0, a, "'\''")
      return a[i]
    }
    function geta(name) {
      split($0, a, "(<driver |/>)")
      split(a[2], b, " ")
      for (item in b) {
        gsub("'\''", "", b[item])
        split(b[item], c, "=")
        key = c[1]
        value = c[2]
        if (key==name) return value
      }
      return ""
    }
    /^ *<name>/ {
      print $3
      bootdev=""
    }
    /^ *<vcpu>/ { printf "\tCores:\t%s\n", $3 }
    /^ *<memory>/ { memory = $3 }
    /^ *<currentMemory>/ {
      current = $3
      if (current==memory)
        printf "\tRAM:\t%1.3f GiB\n", memory/1024/1024
      else
        printf "\tRAM:\t%1.3f / %1.3f GiB\n", current/1024/1024, memory/1024/1024
    }
    /^ *<driver / {
      cache = geta("cache")
      format = geta("type")
    }
    /^ *<source (file|dev)=/ {
      disk_path = get0(2)
    }
    /^ +<target dev=/ {
      target = get0(2)
    }
    /<\/disk>/ {
      printf "\tDisk:\t%-10s %s\n", target, disk_path
      printf "\t\tformat=%s, cache=%s\n", format, cache
    }
    /^ *<interface type=/ {
      interface_type = get0(2)
    }
    /^ *<source .*=/ {
      source = get0(2)
    }
    /^ *<mac address=/ {
      mac_address = get0(2)
    }
    /<\/interface>/ {
      if (interface_type=="bridge") {
        printf "\tBridge:\t%-10s %s\n", source, mac_address
      } else if (interface_type=="network") {
        printf "\tNetwork:\t%-10s %s\n", source, mac_address
      } else {
        printf "\tUNKNOWN:\t%s", mac_address
      }
    }
    /^ *<boot dev=/ {
      if (bootdev) {
        bootdev=bootdev", "get0(2)
      } else {
        bootdev=get0(2)
      }
    }
    /^<\/domain>/ {
      printf "\tBoot:\t%s\n", bootdev
    }
  '
}

diff_guests() {
  g1=`mktemp /tmp/$1,XXXXXXXXXX`
  g2=`mktemp /tmp/$2,XXXXXXXXXX`
  virsh dumpxml $1 >> $g1
  virsh dumpxml $2 >> $g2
  diff -u $g1 $g2
  rm -f $g1 $g2
}

edit_guest() {
  guest="$1"
  g1=`mktemp /tmp/$guest,1,XXXXXXXXXX`
  g2=`mktemp /tmp/$guest,2,XXXXXXXXXX`
  if [ "$2" = "-r" ]; then
    R=$(($3*1048576))
    virsh dumpxml $guest \
      | sed 's|<memory>[0-9]*</memory>|<memory>'$R'</memory>|' \
      | sed 's|<currentMemory>[0-9]*</currentMemory>|<currentMemory>'$R'</currentMemory>|' \
      > $g2
    head $g2
    # set currentMemory for running guest
    virsh setmem $guest $R
  elif [ "$2" = "-c" ]; then
    virsh dumpxml $guest \
      | sed 's|<vcpu>[0-9]*</vcpu>|<vcpu>'$3'</vcpu>|' \
      > $g2
  elif [ "$2" = "-b" ]; then
    shift 2
    bootdevs=""
    for i in "$@"; do
      bootdevs="$bootdevs  <boot dev='$i'/>\n"
    done
    virsh dumpxml $guest \
      | grep -v '^ *<boot dev=' \
      | sed "s|</os>|$bootdevs  </os>|" \
      > $g2
    #grep -e 'os>' -e '<boot' $g2
  else
    virsh dumpxml $guest >> $g1
    ${EDITOR:-vi} $g1
    awk '
      BEGIN {
        FS="<[current/]*[Mm]emory>"
      }
      /<memory>[0-9]+[gG]<\/memory>/ {
        gsub($2,$2*1048576);
        print $0;
        next;
      }
      /<currentMemory>[0-9]+[gG]<\/currentMemory>/ {
        gsub($2,$2*1048576);
        print $0;
        next;
      }
      { print $0 }
    ' < $g1 >> $g2
  fi
  virsh define $g2
  virsh_info info $guest
  rm -f $g1 $g2
}

if [ "${1::2}" == "--" ]; then
  CMD="${1:2}"
else
  CMD="$1"
fi

if [ "$CMD" = "list" -o "$CMD" = "l" ]; then
  virsh_list "$2"
elif [ "$CMD" = "info" -o "$CMD" = "i" ]; then
  virsh_info "$@"
elif [ -z "$2" ]; then
  show_help
elif [ "$CMD" = "restart" -o "$CMD" = "reboot" ]; then
  shutdown_and_restart $2 restart
elif [ "$CMD" = "shutdown" -o "$CMD" = "stop" ]; then
  shutdown_and_restart $2
elif [ "$CMD" = "start" ]; then
  start_guest $2
elif [ "$CMD" = "rename" ]; then
  shift
  rename_guest "$@"
elif [ "$CMD" = "move" ]; then
  move_guest "$@"  
elif [ "$CMD" = "migrate" ]; then
  virsh migrate $4 $2 qemu+ssh://root@$3/system
elif [ "$CMD" = "edit" ]; then
  shift
  edit_guest "$@"
elif [ "$CMD" = "set" ]; then
  GUEST="$3"
  SUBCMD="$2"
  shift 3
  case $SUBCMD in
    ram|memory|mem)
      edit_guest $GUEST -r "$1"
      ;;
    cpu|core|cores)
      edit_guest $GUEST -c "$1"
      ;;
    boot)
      edit_guest $GUEST -b "$@"
      ;;
  esac
elif [ "$CMD" = "console" -o "$CMD" = "con" -o "$CMD" = "c" ]; then
  shift
  virsh console "$@"
elif [ "$CMD" = "diff" ]; then
  diff_guests $2 $3
else
  show_help
fi
