#!/bin/bash

show_help() {
  echo "Usage: $0 [-f] [-d] [-u] [-c] [-a] [-D disk] guest_name [partitions ...]"
  echo "       $0 [-r|--restore] guest filename.gz"
  echo "Example: $0 -a www_server p1 p2"
  echo "         $0 -a www_server 1 2"
  echo "         $0 --restore dir/filename.gz > /dev/..."
  exit
}

if [ -z "$1" ]; then
  show_help
fi

TGT="ftp://vs@backupserver/ImageBackup"
RETURN_STATUS=0

compress="pigz -c"; cext=".gz" # fastest compression
#compress="gzip -c -1"; cext=".gz"
#compress="xz -c -3 -T 24"; cext=".xz"
#compress="bzip2 -c"; cext=".bz2"
#compress="lbzip2 -c"; cext=".bz2"
#compress="cat"; cext="" # no compression

decompress="pigz -d"

checksum="sha256sum"
#checksum="md5deep"

date=`date +%Y%m%d`

# override parameters from configuration file
[ -f /etc/sysconfig/virt-backup ] && . /etc/sysconfig/virt-backup

drop_cache() {
  echo 3 > /proc/sys/vm/drop_caches
}

cycle_dropping_buffers() {
  # drop buffers in intervals to avoid swapping
  trap "exit" USR1
  for i in `seq 300`; do
    drop_cache
    sleep 1m
  done
}

# check for utilities
check_bin() {
  if [ ! -x /usr/bin/$1 ]; then
    if [ ! -x /bin/$1 ]; then
      echo "ERROR: missing command: $1"
      exit 101
    fi
  fi
}
check_bin pv
check_bin $compress
check_bin $checksum
check_bin mktemp
check_bin getopt

check_guest_state() {
  if [ "$FORCE_BACKUP" = "yes" ]; then
    return
  fi
  # exit if domain is already running
  if [ "`virsh domstate $1`" = "running" ]; then
    echo "ERROR: Domain is already running!"
    exit 100
  fi
}

backup_disk() {
  # backup_disk /dev/mapper/diskname [size]
  if [[ $1 == *::* ]]; then
    dev="${1##*::}"
    name="${1%%::*}"
  else
    dev="$1"
    name=`basename $1`
  fi
  cycle_dropping_buffers & DROPPER_PID=$! # run on background and save PID
  echo "Starting backup for $1 ..."
  chksum_file=`mktemp /tmp/virt-backup-sha256.XXXXXXXX`
  if [ "$2" ]; then
    # pv for CentOS 6 does not accept --stop-at-size parameter
    dd if=$dev bs=$2 count=1
  else
    pv $dev
  fi | \
    tee >($checksum > $chksum_file) | \
    $compress | \
    curl -s -S -T - -n $TGT/$guest/${name}_${date}.img$cext
  exit_status=$?
  if [ "$exit_status" != "0" ]; then
    RETURN_STATUS="$exit_status"
  fi
  kill -USR1 $DROPPER_PID
  sleep 1s
  kill -TERM $DROPPER_PID 2> /dev/null
  drop_cache
  cat $chksum_file | curl -s -S -T - -n $TGT/$guest/${name}_${date}.sha256
  rm -f $chksum_file
}

backup_guest() {
  guest=$1
  shift

  # create directory
  curl -s -n \
    -Q "mkdir ImageBackup/$guest" \
    -Q "chmod 0770 ImageBackup/$guest" \
    $TGT #2>/dev/null

  if [ "$VB_DISKS" = "none" ]; then
    disks=""
  elif [ "$VB_DISKS" ]; then
    disks="$VB_DISKS"
  else
    disks=`virsh-xpath --guest $guest --query '//devices/disk/source/@dev' | cut -d'"' -f2`
  fi

  # upload guest configuration
  if [ "$GUEST_CONFIG" = "yes" ]; then
    virsh dumpxml $guest | curl -s -S -T - -n $TGT/$guest/${guest}_${date}.xml
  fi

  # send files
  for disk in $disks; do
    if [ "$1" ]; then
      backup_disk $disk 512 # backup boot sector and partition table
      for part in "$@"; do
        backup_disk "${disk}${part}"
      done
    else
      backup_disk $disk $DISK_SIZE
    fi
  done
}

restore_disk() {
  if [ -z "$2" ]; then
    fn="$1"
  else
    fn="$1/$2"
  fi
  echo "Restore: $TGT/$fn" >&2
  curl -S -n "$TGT/$fn" | $decompress | pv
}

TEMP=$(getopt -o fudcars:D:CNh -l force,up,down,all,auto,console,restore,size:,disk:,config,no-config,help -n virt-backup -- "$@")
if [ $? != 0 ]; then
  echo "Error parsing arguments!" >&2
  exit 1
fi
eval set -- "$TEMP"

DISK_SIZE=""
GUEST_CONFIG="yes"
VB_DISKS=""
while true; do
  case "$1" in
    -f|--force)
        FORCE_BACKUP=yes
        shift
        ;;
    -u|--up)
        POST_START=yes
        shift
        ;;
    -d|--down)
        PRE_STOP=yes
        shift
        ;;
    -c|--console)
        POST_CONSOLE="--console"
        shift
        ;;
    -a|--all|--auto)
        POST_START=yes
        PRE_STOP=yes
        POST_CONSOLE="--console"
        shift
        ;;
    -r|--restore)
        shift 2 # remove "--" too
        restore_disk "$@"
        exit
        ;;
    -D|--disk)
        VB_DISKS="$VB_DISKS $2"
        shift 2
        ;;
    -s|--size)
    	DISK_SIZE="$2"
    	shift 2
    	;;
    -C|--config)
        VB_DISKS="none"
        shift
        ;;
    -N|--no-config)
        GUEST_CONFIG=no
        shift
        ;;
    -h|--help)
        show_help
        ;;
    --)
        shift
        break
        ;;
    \?)	echo "Unknown parameter: $1"
        exit 1
        ;;
  esac  
done

if [ "$PRE_STOP" = "yes" ]; then
  unset FORCE_BACKUP
  virsh shutdown $1
  virsh console $1
fi

check_guest_state "$1"
backup_guest "$@"

if [ "$POST_START" = "yes" ]; then
  virsh start $1 $POST_CONSOLE
fi

exit $RETURN_STATUS