#!/usr/bin/env bash
#
# Vmgr makes deploying templatized virtual machines easy and fast
# Copyright (C)  2016  Aaron Ball <nullspoon@oper.io>
#
# 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/>.
#

#
# Jei! Ég get að skrifa kóðann!!
#

network='default'
prefixmac='52:54:00:7e:70:'
prefixip='192.168.122.'


# Hash to contain direct paths to binaries
# Populated by env_setup
declare -A bins


#
# Standard logging function. Prints standardized output. Timestamps all errors
# and infos.
#
# @param level Log level to print
# @param msg   All remaining arguments to be printed as message
#
function log {
  level=${1}
  shift
  msg=${@}

  d=$(date '+%F %T')

  oldifs=${IFS}
  IFS="
"

  for i in ${msg}; do
    if [[ ${level} == 'error' ]]; then
      echo -n "${d}   ERROR   "
    elif [[ ${level} == 'info' ]]; then
      echo -n "${d}   INFO   "
    elif [[ ${level} == 'debug' ]]; then
      echo -n "${d}   DEBUG   "
    elif [[ ${level} == 'fatal' ]]; then
      echo -n "${d}   FATAL   "
    fi

    echo -e ${i}
  done

  IFS=${oldifs}
  # Exit 1 if message was fatal
  [[ ${1} == 'fatal' ]] && exit 1
}


#
# Searches PATH for all binaries listed as function arguments. Each found
# binary is then put into the global "bins" associative array, using the
# binary's name as the key to its path.
#
# This function also servers to ensure all specified binaries are present and
# in PATH. It will exit code 1 with an error message if any binaries are not
# found.
#
# Note: To use this function, insert "declare -A bins" at top of script).
#
# @param bins All function arguments are the names of binaries to locate.
#
function env_setup {
  local args=${@}

  for i in ${args[@]}; do
    local path=$(which ${i} 2>/dev/null)

    if [[ $? -gt 0 ]]; then
      echo "Could not find binary ${i}. Is it installed?"
      exit 1
    fi

    bins[$i]="${path}"
  done
}


#
# Executes virt-install with the specified arguments.
#
# @param name Name of the virtual machine (what displays in libvirt utils)
# @param disk Path to the disk file
# @param mac  Mac address of the primary nic (eth0)
#             Note that this is important as it will determine the ip address
#             of the system.
#
function install {
  local name=${1}
  local disk=${2}
  local mac=${3}
  local mem=4096
  
  out=$(${bins['virt-install']} \
    --name ${name} \
    --ram=${mem} \
    --virt-type=kvm \
    --disk "${disk}" \
    -w network=default,mac=${mac} \
    --graphics vnc,password=${name} \
    --noautoconsole \
    --import 2>&1)

  # Log the output of virt install
  if [[ $? -gt 0 ]]; then
    log error "${out}"
  else
    log info "${out}"
  fi
}


#
# Clones a disk. Supports two speed options, fast or small.
# Fast just does a simple copy of the disk from source to destination.
# Small uses qemu-img to recompress the disk from source to destination to save
# space.
#
# @param src   Path to source disk to clone
# @param dest  Path to location that the source disk will clone to
# @param speed Speed of conversion (options: small, fast).
#              "fast" uses cp, but is less storage-conscious.
#              "small" uses qemu-img convert, which is slower, but smaller as
#              it recompresses the disk.
#
function clone_disk {
  local src=${1}
  local dest=${2}
  local speed=${3}

  [[ -z ${1} ]] && echo "ERROR: Template path required (arg 1)." && exit 1
  [[ -z ${2} ]] && echo "ERROR: Destination disk path required (arg 2)." && exit 1 
  [[ -z ${3} ]] && echo "ERROR: Clone speed required (arg 3)." && exit 1 

  if [[ ${speed} == 'fast' ]]; then
    ${bins['cp']} ${src} ${dest}
  elif [[ ${speed} == 'small' ]]; then
    ${bins['qemu-img']} convert -O qcow2 -p -c ${src} ${dest}
  else
    log error "Unknown speed: ${speed}"
    exit 1
  fi
}


#
# Reserves an ip on the specified network for mac address ${mac}. Also inserts
# dns entry for the specified name.
#
# @param network Virsh network name of reservation (for virsh net-*)
# @param name    Name of the host (this will resolve in dns)
# @param mac     Mac address to reserve the ip with
# @param ip      IP address to reserve
#
function net_reserve_ip {
  local network=${1}
  local name=${2}
  local mac=${3}
  local ip=${4}

  xml="<host mac='${mac}' name='${name}' ip='${ip}'/>"
  # Perform reservation command
  out=$(${bins['virsh']} net-update ${network} add ip-dhcp-host "${xml}" 2>&1)

  # Log the output of ip reservation operation
  if [[ $? -gt 0 ]]; then
    log error "A problem occured reserving ip address ${ip} for mac ${mac}."
    log error "${out}"
    exit 1
  else
    log info "${out}"
  fi
}


#
# Deletes an ip address reservation for the specified network.
#
# @param network Virsh network name of reservation (for virsh net-*)
# @param name    Name of the host (this is what resolves in dns)
# @param mac     Mac address the IP is reserved to
# @param ip      IP address that will be deleted
#
function net_rm_reservation {
  local network=${1}
  local name=${2}
  local mac=${3}
  local ip=${4}

  xml="<host mac='${mac}' name='${name}' ip='${ip}'/>"
  # Perform reservation command
  ${bins['virsh']} net-update ${network} delete ip-dhcp-host "${xml}"

  if [[ $? -gt 0 ]]; then
    log error "A problem occured deleting ip reservation (${ip}) for mac ${mac}."
    exit 1
  fi
}


#
# Determines the next available vm index based on network mac address usage.
#
# @param network
# @param prefixmac
#
function get_next_index {
  local network=${1}
  local prefixmac=${2}

  # Start at 10, to allow for 0-9 as network-resources that aren't transient.
  index=10

  # Get list of hosts
  hostxml="$(${bins['virsh']} net-dumpxml ${network} | grep '<host' )"

  # For each index, check if the mac address exists
  while [ $(echo ${hostxml} | grep "mac='${prefixmac}${index}'" -c ) -ne 0 ]; do
    # Use this fancy printf statement to ensure we're always zero padded.
    index=$(printf "%0*d" 2 $((${index} + 1)))
  done

  # Unused index found. Echo.
  echo "${index}"
}


#
# Parses a hostname or fqdn to determine a host's 2-digit index (00-99).
#
# @param name Host name or fqdn to parse for host index
#
function get_index_from_name {
  local name=${1}

	# Reset name so we grab whatever is at the front of an fqdn. If name isn't an
  # fqdn, this will return the same.
  name=$(echo ${name} | cut -d '.' -f 1)

  # Get length of the hostname
  local len=$((${#name} - 2))

  # Strip index off end of hostname
  # Note that we assume a 2 digit
  local index=${name:$len:2}

  # Return index
  echo ${index}
}


#
# Creates a new vm from the specified template. Seamlessly handles host index
# increment (appended to prefix). Also handles mac address generation, ip
# reservation and assignment, and disk cloning from template.
#
# @param prefix   Host prefix (index will be appended to this to make hostname)
# @param template Path to disk template the vm will be created from
# @param domain   Domain to be appended to the hostname. Defaults to empty.
#
function vm_new {
  [[ -z ${1} ]] && log fatal "A prefix name for the new vm is required."
  [[ -z ${2} ]] && log fatal "A template disk path is required."

  local prefix=${1}
  local template=${2}
  local domain=${3}

  # Get next mac/index/ip
  local next=$(get_next_index ${network} ${prefixmac})
  local mac="${prefixmac}${next}"
  local ip="${prefixip}${next}"
  local name="${prefix}${next}${domain}"
  local disk="${name}.sda.qcow2"

  # Clone the disk
  log info "Cloning disk from ${template} to ${disk}"
  clone_disk ${template} ${disk} 'fast'

  log info "Reserving ip address ${ip} for mac ${mac}"
  net_reserve_ip ${network} ${name} ${mac} ${ip}

  log info "Installing VM into libvirtd."
  install ${name} ${disk} ${mac} 
}

#
# Shuts down and deletes specified vm Also removes network ip address
# reservation as well as all related storage.
#
# @param name Name of vm to be destroyed
#
function vm_rm {
  [[ -z ${1} ]] && echo "Please a vm name to be deleted." && exit 1

  local name=${1}

  # Get next mac/index/ip/disk
  local index=$(get_index_from_name ${name})
  local mac="${prefixmac}${index}"
  local ip="${prefixip}${index}"
  local disk="${name}.sda.qcow2"

  # Get storage pool name (for refresh later)
  local pool=$(basename $(pwd))

	# Shut the VM down forcibly (it's about to be deleted, so we don't care about
  # a friendly shutdown here)
  ${bins['virsh']} destroy ${name}

	# Refresh the pool after VM destruction. If we don't do this, virsh
	# intermittently fails on storage removal, with a permission denied error.
	# This is because it doesn't know that the disk is no longer in use, so a
  # pool-refresh is required.
  log info "Refreshing storage pool ${pool} to ensure disk removal success."
  poolout=$(${bins['virsh']} pool-refresh ${pool} 2>&1)

  if [[ $? -gt 0 ]]; then
    log error "${poolout}"
  else
    log info "${poolout}"
  fi

  # Delete the VM and remove all of its storage
  log info -n "Removing VM ${name} and all storage."
  undefout=$(${bins['virsh']} undefine --remove-all-storage ${name})

  if [[ $? -gt 0 ]]; then
    log error "${undefout}"
  else
    log info "${undefout}"
  fi

  # Remove its network entry
  log info "Deleting ip reservation for ${name}, mac ${mac}, ip ${ip}"
  net_rm_reservation ${network} ${name} ${mac} ${ip}
}


#
# Lists all VMs currently running, their expected ips, and mac addresses.
#
function vm_ls {
  names=$(${bins['virsh']} list --name)
  echo ${bins['virsh']} list --name

  output='Host  IP  MAC\n---- -- ---\n'
  for i in ${names}; do
    index=$(get_index_from_name ${i})
    output="${output}${i}   ${prefixip}${index}   ${prefixmac}${index}\n"
  done
  echo -e ${output} | column -c 80 -t
}


#
# Everyone loves a main function
#
function main {
  action=${1}
  shift;

  # Set up paths to all binary variables
  env_setup virsh virt-install qemu-img cp

  args=(${@})

  if [[ ${action} == 'new' ]]; then
    vm_new ${args[@]}
  elif [[ ${action} == 'rm' ]]; then
    vm_rm ${args[@]}
  elif [[ ${action} == 'ls' ]]; then
    vm_ls
  elif [[ ${action} == '' ]]; then
    log error "Please specify an action (new, rm, ls)"
    exit 1
  else
    log error "Unknown action: ${action}"
    exit 1
  fi
}

main ${@}