#!/usr/bin/env bash
#
# kvmx virtual machine manager
#
# Copyright (C) 2017 Silvio Rhatto - rhatto at riseup.net
#
# 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 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/>.
#

# Basic parameters
VERSION="0.1.0"
BASENAME="`basename $0`"
DIRNAME="`dirname $0`"
ACTION="$1"
VM="$2"
GLOBAL_USER_CONFIG_FOLDER="$HOME/.config/kvmx"
GLOBAL_USER_CONFIG_FILE="$HOME/.config/kvmxconfig"

# Set application base
function __kvmx_set_app_base {
  local dest
  local base

  # Determine if we are in a local or system-wide install.
  if [ -h "$0" ]; then
    dest="$(readlink $0)"

    # Check again as the caller might be a symlink as well
    if [ -h "$dest" ]; then
      base="`dirname $dest`"
      dest="$(dirname $(readlink $dest))"
    else
      base="`dirname $0`"
      dest="`dirname $dest`"
    fi

    # Deal with relative or absolute links
    if [ "`basename $dest`" == "$dest" ]; then
      export APP_BASE="$base"
    else
      export APP_BASE="$dest"
    fi
  else
    export APP_BASE="`dirname $0`"
  fi
}

# Build a SSH command
function __kvmx_ssh_command {
  SSH_COMMAND="ssh -o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no -o LogLevel=FATAL -o ProxyCommand=none -i $1"
}

# Initialize
function __kvmx_initialize {
  __kvmx_set_app_base

  # Alias to be used in config files
  KVMX_BASE="$APP_BASE"

  if [ "$ACTION" == "init" ] || [ "$ACTION" == "list" ]; then
    return
  fi

  if [ -z "$VM" ]; then
    VM="$(basename `pwd`)"
  fi

  # Default parameters
  PORT="$(($RANDOM + 1024))"
  SSH="$(($PORT + 22))"

  # Initalize
  mkdir -p $GLOBAL_USER_CONFIG_FOLDER

  # Load user config
  if [ -e "$GLOBAL_USER_CONFIG_FILE" ]; then
    source $GLOBAL_USER_CONFIG_FILE
  fi

  # Load and check guest config
  if [ "$ACTION" != "init" ] && [ "$ACTION" != "list" ] && [ "$ACTION" != "ls" ] && [ "$ACTION" != "edit" ] && [ "$ACTION" != "usage" ]; then
    if [ ! -e "$GLOBAL_USER_CONFIG_FOLDER/$VM" ]; then
      echo "$BASENAME: config not found: $GLOBAL_USER_CONFIG_FOLDER/$VM"
      exit 1
    else
      source $GLOBAL_USER_CONFIG_FOLDER/$VM
    fi

    if [ -z "$image" ]; then
      if [ -z "$image_base" ]; then
        image_base="$HOME/.local/share/kvmx"
      fi

      image="$image_base/$VM/box.img"
    fi

    # Box and folder config
    KVMXFILE="`readlink $GLOBAL_USER_CONFIG_FOLDER/$VM`"
    KVMX_PROJECT_FOLDER="`dirname $KVMXFILE`"
    STORAGE="`dirname $image`"
    STATE_DIR="$STORAGE/state/$VM"
    PIDFILE="$STATE_DIR/pid"
    PORTFILE="$STATE_DIR/port"
    SSHFILE="$STATE_DIR/ssh"
    LOGFILE="$STATE_DIR/log"
    SPICEFILE="$STATE_DIR/spice"

    if [ -e "$STORAGE/$VM.key" ]; then
      SSHKEY="$STORAGE/$VM.key"
    else
      SSHKEY="$APP_BASE/share/ssh/insecure_private_key"
    fi

    if [ ! -z "$user" ]; then
      SSH_LOGIN="$user"
    else
      SSH_LOGIN="user"
    fi

    __kvmx_ssh_command $SSHKEY

    mkdir -p $STATE_DIR

    if [ ! -e "$image" ] && [ "$ACTION" != "up" ] && [ "$ACTION" != "purge" ] && [ "$ACTION" != "destroy" ]; then
      echo "$BASENAME: file not found: $image"
      exit 1
    fi
  fi
}

# Run spice client
function kvmx_spice {
  if ! kvmx_running; then
    echo "$BASENAME: guest $VM is not running"
    exit 1
  fi

  # Ensure we have the right port configuration: we can also be
  # running directly from command line.
  PORT="`cat $PORTFILE`"

  if [ -z "$PORT" ]; then
    echo "$BASENAME: cannot get spice port for $VM."
    exit 1
  fi

  # https://lists.freedesktop.org/archives/spice-devel/2013-September/014643.html
  SPICE_NOGRAB=1 spicec --host localhost --port $PORT &> $SPICEFILE &
  #spicy -h localhost -p $PORT
  #remote-viewer spice://localhost:$PORT

  # Give time to connect
  sleep 5

  # Fix window titles
  if which /usr/bin/xdotool &> /dev/null; then
    xdotool search --name "SPICEc:0" set_window --name $VM
  fi
}

# Bring virtual machine up
function kvmx_up {
  if kvmx_running; then
    echo "$BASENAME: guest $VM is already running"
    exit 1
  fi

  if [ ! -z "$shared_folder" ]; then
    # Get absolute path of shared folder relative to project path
    shared_folder="`cd $KVMX_PROJECT_FOLDER && cd $shared_folder &> /dev/null && pwd`"
    local shared="-fsdev local,id=shared,path=$shared_folder,security_model=none -device virtio-9p-pci,fsdev=shared,mount_tag=shared"
  fi

  if [ ! -z "$port_mapping" ]; then
    local hostfwd=",$port_mapping"
  fi

  # Check if image exists, create otherwise
  if [ ! -e "$image" ]; then
    if [ ! -z "$basebox" ]; then
      if [ -e "$GLOBAL_USER_CONFIG_FOLDER/$basebox" ]; then
        baseimage="`kvmx list_image $basebox`"
        basekey="`basename $baseimage .img`.key"

        if [ ! -e "$baseimage" ]; then
          echo "$BASENAME: could not find basebox $baseimage. Please create it first."
          exit 1
        fi

        echo "Copying base image $baseimage to $image..."
        cp $baseimage $image

        if [ -e "$basekey" ]; then
          imagekey="`basename $image .img`.key"
          cp $basekey     $imagekey
          cp $basekey.pub $imagekey.pub

          # Re-evaluate this if there's a custom SSH key.
          __kvmx_ssh_command $basekey
        fi
      fi
    else
      kvmx-create $GLOBAL_USER_CONFIG_FOLDER/$VM
    fi
  fi

  # Run virtual machine
  # See https://en.wikipedia.org/wiki/Nohup#Overcoming_hanging
  nohup kvm -m 2048 -name $VM -drive file=$image,if=virtio -vga qxl $shared \
      -spice port=$PORT,addr=127.0.0.1,disable-ticketing,streaming-video=off,jpeg-wan-compression=never,playback-compression=off,zlib-glz-wan-compression=never,image-compression=off \
      -device virtio-serial-pci \
      -device virtserialport,chardev=spicechannel0,name=com.redhat.spice.0 \
      -chardev spicevmc,id=spicechannel0,name=vdagent \
      -smp 2 -soundhw ac97 -cpu host -balloon virtio  \
      -net nic,model=virtio \
      -net user,hostfwd=tcp:127.0.0.1:$SSH-:22$hostfwd &> $LOGFILE < /dev/null &

  PID="$!"

  # Save state
  echo $PID  > $PIDFILE
  echo $PORT > $PORTFILE
  echo $SSH  > $SSHFILE

  if [ "$run_spice_client" == "1" ]; then
    kvmx_spice
  fi

  let ssh_attempts="0"
  echo -n "Waiting for machine to boot..."
  while true; do
    echo true | $SSH_COMMAND -o ConnectTimeout=2 -o NumberOfPasswordPrompts=0 -p $SSH $SSH_LOGIN@127.0.0.1 &> /dev/null && break
    echo -n "."
    let ssh_attempts++

    if [ "$ssh_attempts" == "10" ]; then
      echo "$BASENAME: timeout or access denied when trying to SSH into $VM."
      echo "$BASENAME: please check if the image is in a good state and if it accepts passwordless ssh connections using kvmx insecure privkey"
      kvmx_stop
      exit 1
    fi

    sleep 2
  done
  echo " done."
  #sleep 5
  #echo ""

  # Somehow it is starting before DBUS and then crashing, so we try to start again
  echo "Ensure spice-vdagent is running..."
  echo "sudo /usr/sbin/service spice-vdagent start" | kvmx_ssh

  kvmx_hostname

  if [ ! -z "$shared_folder" ] && [ ! -z "$shared_folder_mountpoint" ]; then
    echo "Mounting $shared_folder on $shared_folder_mountpoint on guest..."
    echo "sudo mkdir -p $shared_folder_mountpoint" | kvmx_ssh
    echo "sudo mount -t 9p -o trans=virtio shared $shared_folder_mountpoint -oversion=9p2000.L,posixacl,cache=loose" | kvmx_ssh
  fi

  kvmx_status
}

# Set hostname
function kvmx_hostname {
  if ! kvmx_running; then
    echo "$BASENAME: guest $VM is not running"
    exit 1
  fi

  echo "Setting hostname..."
  $SSH_COMMAND -o ConnectTimeout=2 -p $SSH $SSH_LOGIN@127.0.0.1 <<EOF
  ##### BEGIN REMOTE SCRIPT #####
  OLD_HOST="\$(hostname)"

  # Set hosts entry
  if ! grep -q "^127.0.0.1 $hostname.$domain $hostname$" /etc/hosts; then
    echo "127.0.0.1 $hostname.$domain $hostname" | sudo tee -a /etc/hosts > /dev/null
  fi

  echo "$hostname.$domain" | sudo tee /etc/hostname > /dev/null
  sudo hostname $hostname.$domain

  # Remove old hostname from hosts file
  if [ "\$OLD_HOST" != "$hostname.$domain" ]; then
    if grep -q \$OLD_HOST /etc/hosts; then
      sudo sed -i -e "/\$OLD_HOST/d" /etc/hosts
    fi
  fi
  ##### END REMOTE SCRIPT #######
EOF
}

# Display usage
function kvmx_usage {
  echo "$BASENAME $VERSION - virtual machine manager"
  echo ""
  echo "usage: $BASENAME <action> [options]"
  echo ""
  echo "available actions:"
  echo ""
  grep "^function kvmx_" $0 | cut -d ' ' -f 2 | sed -e 's/kvmx_/\t/' | sort
  echo ""
  echo "examples:"
  echo ""
  echo -e "\t$BASENAME list"
  echo -e "\t$BASENAME init  <machine>    [folder]"
  echo -e "\t$BASENAME clone <orig-guest> <dest-folder>"
  echo ""

  local list="`kvmx_list | xargs`"

  if [ ! -z "$list" ]; then
    echo "available virtual machines:"
    echo ""
    echo -e "\t$list"
    echo ""
  fi

  exit 1
}

# Log into the guest using SSH
function kvmx_ssh {
  if ! kvmx_running; then
    echo "$BASENAME: guest $VM is not running"
    exit 1
  fi

  if [ "$ssh_support" != "y" ]; then
    echo "$BASENAME: SSH support for $VM is disabled"
    exit 1
  fi

  # Shift params according to how the program was called:
  # either "kvmx ssh" or "kvmx ssh guest".
  if [ "$ACTION" == "ssh" ]; then
    if [ ! -z "$2" ]; then
      shift 2
    else
      shift 1
    fi
  fi

  SSH="`cat $SSHFILE`"
  $SSH_COMMAND -p $SSH $SSH_LOGIN@127.0.0.1 $*
}

# Suspend the virtual machine
function kvmx_suspend {
  if ! kvmx_running; then
    echo "$BASENAME: guest $VM is not running"
    exit 1
  fi

  PID="`cat $PIDFILE`"
  kill -STOP $PID
}

# Check if a guest is running
function kvmx_running {
  if [ ! -e "$PIDFILE" ]; then
    return 1
  fi

  PID="`cat $PIDFILE`"

  if [ -z "$PID" ]; then
    return 1
  fi

  ps $PID &> /dev/null

  return $?
}

# Resume the guest
function kvmx_resume {
  if ! kvmx_running; then
    echo "$BASENAME: guest $VM is not running"
    exit 1
  fi

  PID="`cat $PIDFILE`"
  kill -CONT $PID
}

# Poweroff the guest
function kvmx_poweroff {
  if ! kvmx_running; then
    echo "$BASENAME: guest $VM is not running"
    exit 1
  fi

  echo /usr/bin/sudo poweroff | kvmx_ssh &> /dev/null
  sleep 3
  kvmx_status
}

# Alias for poweroff
function kvmx_down {
  kvmx_poweroff
}

# Hibernate
function kvmx_hibernate {
  if ! kvmx_running; then
    echo "$BASENAME: guest $VM is not running"
    exit 1
  fi

  echo "which s2disk &> /dev/null && /usr/bin/sudo s2disk" | kvmx_ssh &> /dev/null

  echo "Checking if hibernation was successful..."
  sleep 3

  if kvmx_running; then
    echo "Unable to hibernate guest: please check guest configuration"
    exit 1
  else
    kvmx_status
  fi
}

# Reboot the guest
function kvmx_reboot {
  if ! kvmx_running; then
    echo "$BASENAME: guest $VM is not running"
    exit 1
  fi

  echo /usr/bin/sudo reboot | kvmx_ssh &> /dev/null
  sleep 3
  kvmx_status
}

# Rsync files to the guest
function kvmx_rsync {
  if ! kvmx_running; then
    echo "$BASENAME: guest $VM is not running"
    exit 1
  fi

  ORIG="$3"
  DEST="$4"
  SSH="`cat $SSHFILE`"
  rsync -av -e "$SSH_COMMAND -o Port=$SSH" --rsync-path "sudo rsync" $ORIG/ $SSH_LOGIN@127.0.0.1:$DEST/
}

# List guests
function kvmx_list {
  if [ -e "$GLOBAL_USER_CONFIG_FOLDER" ]; then
    ls $GLOBAL_USER_CONFIG_FOLDER
  fi
}

# Alias to list command
function kvmx_ls {
  kvmx_list
}

# Upgrade guest
function kvmx_upgrade {
  echo "sudo apt-get update && sudo apt-get dist-upgrade -y && sudo apt-get autoremove -y" | kvmx_ssh
}

# Initializes a new guest
function kvmx_init {
  FOLDER="$3"

  if [ -z "$FOLDER" ]; then
    if [ -z "$VM" ]; then
      VM="$(basename `pwd`)"
      FOLDER="$(dirname `pwd`)/$VM"
    else
      FOLDER="$(pwd)/$VM"
    fi
  fi

  if [ -e "$GLOBAL_USER_CONFIG_FOLDER/$VM" ]; then
    echo "$BASENAME: guest $VM already exists"
    exit 1
  fi

  if [ ! -d "$FOLDER" ]; then
    mkdir -p $FOLDER
  fi

  # Ensure we have an absolute folder name
  FOLDER="`cd $FOLDER &> /dev/null && pwd`"

  # Copy config from template
  if [ ! -e "$FOLDER/kvmxfile" ]; then
    cp $APP_BASE/kvmxfile $FOLDER/
    sed -i -e "s|hostname=\"machine\"|hostname=\"$VM\"|g" $FOLDER/kvmxfile
  fi

  # Create config entry
  ( cd $GLOBAL_USER_CONFIG_FOLDER && ln -s $FOLDER/kvmxfile $VM )
}

# Clone a guest
function kvmx_clone {
  if kvmx_running; then
    echo "$BASENAME: orig $VM is running, cannot clone."
    exit 1
  fi

  FOLDER="$3"
  DEST="`basename $FOLDER`"

  if [ -z "$FOLDER" ]; then
    kvmx_usage
  fi

  # Check if dest machine exists
  if [ -e "$GLOBAL_USER_CONFIG_FOLDER/$DEST" ]; then
    echo "$BASENAME: destination guest $DEST already exists."
    exit 1
  fi

  if [ -d "$FOLDER" ]; then
    echo "$BASENAME: destination $FOLDER already exists."
    exit 1
  fi

  # Ensure we have an absolute folder name
  mkdir -p $FOLDER
  FOLDER="`cd $FOLDER &> /dev/null && pwd`"
  rmdir $FOLDER

  # Copy image and configuration
  cp -r `dirname $image` $FOLDER/
  ( cd $GLOBAL_USER_CONFIG_FOLDER && ln -s $FOLDER/kvmxfile $DEST )

  # Update config file
  new_image="$FOLDER/`basename $image`"
  sed -i -e "s|image=\"$image\"|image=\"$new_image\"|g" $GLOBAL_USER_CONFIG_FOLDER/$DEST
  sed -i -e "s|hostname=\"$VM\"|hostname=\"$DEST\"|g"   $GLOBAL_USER_CONFIG_FOLDER/$DEST
}

# Edit guest config
function kvmx_edit {
  if [ -z "$EDITOR" ]; then
    EDITOR="vi"
  fi

  if [ -e "$GLOBAL_USER_CONFIG_FOLDER/$VM" ]; then
    $EDITOR $GLOBAL_USER_CONFIG_FOLDER/$VM
  else
    echo "$BASENAME: $GLOBAL_USER_CONFIG_FOLDER/$VM: file not found."
  fi
}

# Stop a guest
function kvmx_stop {
  if kvmx_running; then
    PID="`cat $PIDFILE`"
    kill $PID
  fi
}

# Destroy a guest
function kvmx_destroy {
  kvmx_stop

  rm -f  $image
  rm -rf $STATE_DIR

  echo "$BASENAME: removed image and state files, but not the whole`dirname $image` folder."
}

# Shred a guest
function kvmx_shred {
  kvmx_stop

  if which shred &> /dev/null; then
    shred $image
    rm -f $image
  else
    echo "$BASENAME: error shreding $image: shred program not available."
    exit 1
  fi
}

# Wipe a guest
function kvmx_wipe {
  kvmx_stop

  if which wipe &> /dev/null; then
    wipe -f $image
    rm   -f $image
  else
    echo "$BASENAME: error wipeing $image: wipe program not available."
    exit 1
  fi
}

# Purge a guest and all its configuration
function kvmx_purge {
  kvmx_destroy
  rm -f $GLOBAL_USER_CONFIG_FOLDER/$VM
  echo "$BASENAME: removed $GLOBAL_USER_CONFIG_FOLDER/$VM config."
}

# Provision a machine
function kvmx_provision {
  if ! kvmx_running; then
    echo "$BASENAME: guest $VM is not running"
    exit 1
  fi

  if [ -z "$provision_command" ]; then
    echo "$BASENAME: error: parameter provision_command is not configured for $VM."
    exit 1
  fi

  if [ ! -z "$provision_rsync" ]; then
    SSH="`cat $SSHFILE`"
    ORIG="`echo $provision_rsync | cut -d ' ' -f 1`"
    DEST="`echo $provision_rsync | cut -d ' ' -f 2`"

    echo "Syncing provision files into the guest..."
    echo "sudo mkdir -p `dirname $DEST`" | kvmx_ssh
    rsync -av -e "$SSH_COMMAND -o Port=$SSH" --rsync-path "sudo rsync" $ORIG/ $SSH_LOGIN@127.0.0.1:$DEST/
  fi

  echo "Running provision command inside the guest..."
  echo "$provision_command $hostname $domain $mirror" | kvmx_ssh
}

# Print guest image file name
function kvmx_list_image {
  echo $image
}

# Print guest status
function kvmx_status {
  if kvmx_running; then
    echo "$BASENAME: $VM guest is running"
    PID="`cat $PIDFILE`"
    ps $PID
  else
    echo "$BASENAME: $VM guest is stopped"
  fi
}

# Print guest log
function kvmx_log {
  local logs=""

  if [ -s "$LOGFILE" ]; then
    logs="$logs $LOGFILE"
  fi

  if [ -s "$SPICEFILE" ]; then
    logs="$logs $SPICEFILE"
  fi

  tail -F $logs
}

# Rotate SSH keys
function kvmx_rotate_sshkeys {
  # Generate new keypair
  SSHKEY="$STORAGE/$VM.key"
  $DIRNAME/kvmx-keygen $SSHKEY.new "$user@`basename $image .img`"

  # Replace pubkey on server
  echo "touch ~/.ssh/authorized_keys.new && chmod 600 ~/.ssh/authorized_keys.new" | kvmx_ssh
  cat $SSHKEY.new.pub | kvmx_ssh "tee ~/.ssh/authorized_keys.new &> /dev/null"
  echo "mv ~/.ssh/authorized_keys.new ~/.ssh/authorized_keys" | kvmx_ssh

  # Replace keypair locally
  mv $SSHKEY.new     $SSHKEY
  mv $SSHKEY.new.pub $SSHKEY.pub
}

# Dispatch
if type kvmx_$ACTION 2> /dev/null | grep -q 'function'; then
  __kvmx_initialize
  kvmx_$ACTION $*
else
  kvmx_usage
fi