#! /bin/bash

#Usage: /usr/bin/qwdctl controls the availability of virtual machines for the qemu-web-desktop/DARTS service.
#
#  Each entry in the configuration file `/etc/qemu-web-desktop/machines.conf` 
#  spans on 3 lines:
#
#  -  [name.ext] 
#  -  url=[URL to ISO, QCOW2, VDI, VMDK, RAW, VHD/VHDX, QED virtual machine disk, optional]
#  -  description=[description to be shown in the service page] 
#
#  Images listed in the configuration file without a `url=` parameter are
#  expected to be downloaded by hand and installed into
#  `/var/lib/qemu-web-desktop/machines` by the local administrator. Then, just 
#  specify the [name.ext] and description.
# 
# use: 'qwdctl help' to get help about usage and commands/options.

set -e

bin=$(dirname $0)
p=$(basename $0)

F=/opt/homebrew
if [ -d  $F ]; then
  PREFIX=$F
fi

# files to process
F=$PREFIX/etc/qemu-web-desktop/machines.conf
machine_conf=$F
if [ ! -f $F ]; then
  echo "$p: WARNING: can not find QWD machine.conf file ($F)."
fi

F=$PREFIX/usr/share/qemu-web-desktop/html/desktop/index.html
index_html=$F
if [ -f      $PREFIX/var/www/qemu-web-desktop/html/desktop/index.html ]; then
  index_html=$PREFIX/var/www/qemu-web-desktop/html/desktop/index.html
elif [ ! -f  $F ]; then
  echo "$p: WARNING: can not find QWD index.html file ($F)."
fi

# set to yes if apache/httpd include_module does not properly show list in index.html
machines_insert=no

# generated files, should be linked into /usr/share/qemu-web-desktop/html/desktop/
F=$PREFIX/var/lib/qemu-web-desktop
qwdprefix=$F
if [ ! -d $F ]; then
  echo "$p: WARNING: can not find QWD machines directory ($F)."
fi
machine_html=$qwdprefix/machines.html

# search for local config
config_pl=$(realpath "$bin/../config.pl")
if [ -f $config_pl ]; then
  # using local configuration
  true
else
  F=$PREFIX/etc/qemu-web-desktop/config.pl
  config_pl=$F
  if [ ! -f $F ]; then
    echo "$p: WARNING: can not find QWD config.pl file ($F)."
  fi
fi

# search for local cgi/.pl script
cgi=$(realpath "$bin/../cgi-bin")
if [ -x "$cgi/qemu-web-desktop.pl" ]; then
  # local (git)
  cgi="$cgi/qemu-web-desktop.pl"
elif [ -x "$PREFIX/srv/http/cgi-bin/qemu-web-desktop.pl" ]; then
  # Arch/httpd default
  cgi=$PREFIX/srv/http/cgi-bin/qemu-web-desktop.pl
elif [ -x "$PREFIX/var/www/cgi-bin/qemu-web-desktop.pl" ]; then
  # RedHat/httpd and Brew default
  cgi=$PREFIX/var/www/cgi-bin/qemu-web-desktop.pl
else
  # Debian default
  cgi=$PREFIX/usr/lib/cgi-bin/qemu-web-desktop.pl
fi

# identify type of system
if [ -e $PREFIX/etc/initramfs-tools/modules ]; then
	# Debian-class
	DIST=Debian
else
	# Arch-class
	DIST=Arch
fi

# find machines: check for confget/crudini
if   command -v confget 2>&1 >/dev/null; then
  parse_ini=confget
  sections=$(confget -f $machine_conf -q sections);
elif command -v crudini 2>&1 >/dev/null; then
  parse_ini=crudini
  sections=$(crudini --get $machine_conf);
else
  case "$1" in
    download|update|--download|refresh|--refresh)
      echo "$p: $1: WARNING: confget/crudini not found. Install any of these."
    ;;
  esac
fi

case "$1" in
# ------------------------------------------------------------------------------
    download|update|--download)
	mkdir -p $qwdprefix/machines || true
	cd $qwdprefix/machines
	
	for i in $sections ; do
	    mkdir -p downloads/$i || true
	    if [ "x$parse_ini" = "xconfget" ]; then
	      u=$(confget -f $machine_conf -s $i url)
	    else
	      u=$(crudini --get $machine_conf $i url || true)
	    fi
	    if [ "$u" ] ; then
		cd downloads/$i
		echo "Getting $u"
		wget -N $u
		cd ../..
	    fi
	    vm=$(ls -t downloads/$i/* | head -1)
	    if [ -e "$vm" ] ; then
		ln -sf $vm $i
	    fi
	done
	$p refresh
	;;
# ------------------------------------------------------------------------------
    refresh|--refresh)
	mkdir -p $qwdprefix/snapshots || true
	if getent passwd _qemu-web-desktop > /dev/null 2>&1; then
	  chown _qemu-web-desktop $qwdprefix/snapshots # when using MPM ITK
	fi
	mkdir -p $qwdprefix/machines || true
	cd $qwdprefix/machines
	echo "Updating $machine_conf"
	echo "      -> $qwdprefix/machines"
	
	# list of machines
	t=$(mktemp $machine_html.XXXXXX)
	chmod 644 $t
	chmod a+rx $qwdprefix
	for i in $sections; do
	    if [ "x$parse_ini" = "xconfget" ]; then
	      d=$(confget -f $machine_conf -s $i description)
	    else
	      d=$(crudini --get $machine_conf $i description)
	    fi
	    if [ -e $i ] ; then
	    	    if [ "$d" ]; then
		    	# add entry when VM file and descr are given
		    	echo "Found $i '$d'"
			echo "<option value='$i'>$d</option>" >> $t
		    fi
	    fi
	done
	mv $t $machine_html
	# handle case when apache mod_include does not work...
	if [ "x$machines_insert" = "xyes" ]; then
		t=$(mktemp $index_html.XXXXXX)
		lead="<\!--BEGIN_MACHINE_LIST-->"
		tail="<\!--END_MACHINE_LIST-->"
		sed -e "/$lead/,/$tail/{ /$lead/{p; r $machine_html
        }; /$tail/p; d }" $index_html > $t
        mv $t $index_html
		chmod a+r $index_html
	fi
	;;
# ------------------------------------------------------------------------------
	  status|--status)
	echo "$p Configuration"
	echo "  config:   $config_pl"
	echo "  machines: $machine_conf"
	echo "            $qwdprefix/machines/"
	echo "  landing:  $index_html"
	echo " "
	echo "Available machines:"
	echo "$sections"
	echo " "
	echo "Active sessions:"
	echo "session_ID:user:VM            | #cpu | #mem[MB]"
	echo "------------------------------|------|---------"
	t=$(ps aux | grep qemu-web-desktop)
	name=$(echo "$t" | grep -oP '(?<=\-name )[^ ]*' )
	cpu=$(echo "$t"  | grep -oP '(?<=\-smp )[^ ]*' )
	mem=$(echo "$t"  | grep -oP '(?<=\-m )[^ ]*' )
	table=$(printf '%s\n' "$name" "$cpu" "$mem" | pr -3 -T)
	u=$(echo "$table" | uniq )
	echo "$u"
	;;
# ------------------------------------------------------------------------------
	  start|launch|--start)
	# 2nd arg it the VM
	if [ "x$2" = "x" ]; then
	  echo "Usage: $p start VM [options]"
	  echo "Configuration: $config_pl"
	  echo "Script:        $cgi"
	  echo "Availabe options are:"
	  $cgi -h
	  exit 1
	fi
	DIR=$(dirname $2)
	DIR=$(realpath $DIR)
	VM=$(basename $2)
	shift 2
	# snapshot_alloc_cpu snapshot_alloc_mem
	echo "Launching $DIR/$VM $@"
	echo "Connect to (log in /tmp/${VM}.log):"
	$cgi --dir_snapshots=/tmp --dir_cfg=/tmp --dir_machines=$DIR --machine=$VM --oneshot=1 --service_max_mem_fraction_nb_per_user=0.8 --service_max_cpu_fraction_nb_per_user=0.8 --snapshot_alloc_disk=30 $@ 2>&1 | tee /tmp/${VM}.log | awk '/URL:/ { print $4 }'
	echo "Cleaning remaining sessions..."
	$cgi --service_purge=1 --dir_snapshots=/tmp --dir_cfg=/tmp 2>&1 | grep -i "stop"
	echo "Done $VM"
	;;
# ------------------------------------------------------------------------------
	  stop|--stop)
	if [ "x$2" = "x" ]; then
	  echo "Stopping all local VMs"
	  $cgi --service_purge=1 --dir_snapshots=/tmp --dir_cfg=/tmp 2>&1 | grep -i "stop"
	  exit 0
	fi
	t=$(ps aux | grep qemu-web-desktop | grep $2)
	name=$(echo "$t" | grep -oP '(?<=\-name )[^ ]*' )
	pid=$( echo "$t" | awk '{ print $2 }' )
	table=$(printf '%s\n' "$name" | pr -1 -Ts'\t')
	u=$(echo "$table" | uniq )
	echo "Stopping:"
	echo "$u"
	kill $pid
	;; 
	
# ------------------------------------------------------------------------------
	  gpu_reattach|gpu_clean|gpu_unlock|--gpu_unlock)
	
	echo "Cleaning existing GPU virtualization (remove VFIO pass-through)."
	echo "All detached/locked GPUs will be returned to the server."
	echo "The following files will be modified:"
	echo "  /etc/default/grub      (GRUB_CMDLINE_LINUX_DEFAULT) update"
	echo "  /etc/security/limits.conf                           kept"
	echo "  /etc/modprobe.d/vfio.conf                           remove"
	echo "  /etc/udev/rules.d/10-qemu-hw-users.rules            remove"
	if [ "x$DIST" = "xDebian" ]; then
	echo "  /etc/initramfs-tools/modules                        update"
	echo "  /etc/systemd/system/apache2.service.d/override.conf remove"
	else
	echo "  /etc/mkinitcpio.conf                                update"
	echo "  /etc/systemd/system/httpd.service.d/override.conf   remove"
	echo "  NOTE: require package 'mkinitcpio'"
	fi
	read -p "Continue (y/N)?" choice
	case "$choice" in 
		y|Y ) echo "Proceeding";;
		*) echo "Aborting." && exit 1;;
	esac
	
	# restore default GRUB options
	FILE=$PREFIX/etc/default/grub
	DATE=$(date +"%Y-%m-%d-%H-%M")
	echo "Updating  $FILE (old stored as $FILE.$DATE)"
	cp $FILE $FILE.$DATE
	sed -i "s/^GRUB_CMDLINE_LINUX_DEFAULT/#GRUB_CMDLINE_LINUX_DEFAULT/g" $FILE
	echo "GRUB_CMDLINE_LINUX_DEFAULT=\"quiet splash\"" >> $FILE
	
	# comment modules
	if [ "x$DIST" = "xDebian" ]; then
	  FILE=$PREFIX/etc/initramfs-tools/modules
	else
	  if ! command -v mkinitcpio 2>&1 >/dev/null; then
	    echo "$p: ERROR: Missing mkinitcpio to proceed. Install it, and then retry."
	    exit 1
	  fi
	  FILE=$PREFIX/etc/mkinitcpio.conf
	fi
	if [ -e $FILE ]; then
		echo "Updating  $FILE (old stored as $FILE.$DATE)"
		cp $FILE $FILE.$DATE
		if [ "x$DIST" = "xDebian" ]; then
		sed -i "s/^vfio/#vfio/g"                 $FILE
		sed -i "s/^vhost-netdev/#vhost-netdev/g" $FILE
		else
		sed -i "s/^MODULES+=/#MODULES+=/g"       $FILE
		sed -i "s/^HOOKS+=/#HOOKS+=/g"           $FILE
		fi
	fi

	rm -f $PREFIX/etc/modprobe.d/vfio.conf
	rm -f $PREFIX/etc/udev/rules.d/10-qemu-hw-users.rules
	if [ "x$DIST" = "xDebian" ]; then
	rm -f $PREFIX/etc/systemd/system/apache2.service.d/override.conf
	else
	rm -f $PREFIX/etc/systemd/system/httpd.service.d/override.conf
	fi
	
	if [ "x$DIST" = "xDebian" ]; then
		update-initramfs -u
		update-grub
	else
		mkinitcpio -p linux
		grub-mkconfig -o $PREFIX/boot/grub/grub.cfg  || echo "WARNING: Failed to update GRUB"
	fi
	udevadm control --reload-rules
	udevadm trigger
	systemctl daemon-reload
	
	# request to uncomment GPU section in index.html
	echo "-----------------------------------------------------------------"
	echo "Check the above files, then COMMENT the GPU section in the file"
	echo "  $index_html"
	echo "You may use 'sudo $p edit web' to do so."
	echo "Now please reboot to activate changes."
	echo "-----------------------------------------------------------------"
	;;
	
# ------------------------------------------------------------------------------
	  gpu|gpu_lock|--gpu)
	GPU_ids=$2
  	
  	# get CPU model
	CPU_type=
	grep -qi "intel"     /proc/cpuinfo && CPU_type=intel_iommu
	grep -qi "amd"       /proc/cpuinfo && CPU_type=amd_iommu
	if [ "x$CPU_type" = "x" ]; then
		echo "ERROR: GPU support can only be configured for AMD and Intel CPUs."
		exit 1
	fi
  	
  	# list available GPUs
	echo "Available GPUs. The GPU_ids are usually shown as [vendor:model]"
	GPU_avail=`lspci -nnv | grep -i "VGA\|3d controller"`
	echo "$GPU_avail"
	GPU_nb=$(echo "$GPU_avail" | wc -l)
	if [ "$GPU_nb" -lt "2" ]; then
		echo "ERROR: you need at least 2 GPU models, and keep one for your display."
		exit 1
	fi
	echo "NOTE: keep one GPU for your current display."
	if [ "x$GPU_ids" = "x" ]; then
		read -p "Enter vendor:model ID e.g. 10de:1d01 (Ctrl-C to abort): " GPU_ids
	fi
	if [ "x$GPU_ids" = "x" ]; then
		echo "Usage: $p gpu GPU_ids";
		echo "  The GPU_ids should be e.g. '10de:1d01' or '10de:1d01,10de:1d01'"
		exit 1
	fi

	GPU_firstID=`echo $GPU_ids | cut -d ',' -f1`
	echo $GPU_avail | grep -q "$GPU_firstID" || GPU_avail=no
	if [ "x$GPU_avail" = "xno" ]; then
		echo "ERROR: GPU $GPU_ids model not available."
		exit 1
	fi

	# display message and wait for confirmation
	echo "Ready to configure GPU $GPU_ids for $CPU_type (set VFIO pass-through)."
	echo "All these GPU models will be detached from the server and made usable for virtualization."
	echo "WARNING: Make sure you keep a display for your system."
	echo "The following files will be modified ($DIST):"
	echo "  /etc/default/grub      (GRUB_CMDLINE_LINUX_DEFAULT) update"
	echo "  /etc/security/limits.conf                           append"
	echo "  /etc/modprobe.d/vfio.conf                           create"
	echo "  /etc/udev/rules.d/10-qemu-hw-users.rules            create"
	if [ "x$DIST" = "xDebian" ]; then
	echo "  /etc/initramfs-tools/modules                        append"
	echo "  /etc/systemd/system/apache2.service.d/override.conf create"
	else
	echo "  /etc/mkinitcpio.conf                                append"
	echo "  /etc/systemd/system/httpd.service.d/override.conf   create"
	echo "  NOTE: require package 'mkinitcpio'"
	fi
	read -p "Continue (y/N)?" choice
	case "$choice" in 
		y|Y ) echo "Proceeding";;
		*) echo "Aborting." && exit 1;;
	esac

	# GRUB
	FILE=$PREFIX/etc/default/grub
	grep -i "^GRUB_CMDLINE_LINUX_DEFAULT" $FILE | grep "_iommu" && echo "ERROR: $FILE seems to already contain IOMMU keyword. Please check content and/or clean above files. You may use: $p gpu_unlock" && exit 1
	DATE=$(date +"%Y-%m-%d-%H-%M")
	echo "Updating  $FILE (old stored as $FILE.$DATE)"
	cp $FILE $FILE.$DATE
	sed -i "s/^GRUB_CMDLINE_LINUX_DEFAULT=\"/GRUB_CMDLINE_LINUX_DEFAULT=\"$CPU_type=on iommu=pt vfio-pci.ids=$GPU_ids /g" $FILE

	# VFIO modules
	FILE=$PREFIX/etc/modprobe.d/vfio.conf
	echo "Creating  $FILE"
	dd status=none of=${FILE} << EOF
# $FILE
options vfio-pci ids=$GPU_ids disable_vga=1
EOF
	
	# Initramfs modules
	if [ "x$DIST" = "xDebian" ]; then
		FILE=$PREFIX/etc/initramfs-tools/modules
		MODULES=$(cat << EOF
# $FILE for qemu-web-desktop
vfio
vfio_iommu_type1
vfio_pci
vfio_virqfd
vhost-netdev
EOF
)
	else
		FILE=$PREFIX/etc/mkinitcpio.conf
		if ! command -v mkinitcpio 2>&1 >/dev/null; then
	    echo "$p: ERROR: Missing mkinitcpio. Install it, and then retry."
	    exit 1
	  fi
		MODULES=$(cat << EOF
# $FILE for qemu-web-desktop
MODULES+=(vfio vfio_iommu_type1 vfio_pci)
HOOKS+=(modconf)
EOF
)
	fi
	# append only if token not found in FILE
	echo "Appending $FILE"
	if [ -e $FILE  ]; then cp $FILE $FILE.$DATE; fi
	# append only if token not found in FILE
	grep -q "vfio_pci" "$FILE" || echo "$MODULES" >> "$FILE"

	# udev Rules
	FILE=$PREFIX/etc/udev/rules.d/10-qemu-hw-users.rules
	if [ -e $FILE ]; then
		echo "WARNING:  $FILE is kept (already there)."
	else
		echo "Creating  $FILE"
		dd status=none of=${FILE} << EOF
# $FILE for qemu-web-desktop
SUBSYSTEM=="vfio", OWNER="root", GROUP="kvm"
EOF
	fi
  
	# limits: append at end
	FILE=$PREFIX/etc/security/limits.conf
	echo "Appending $FILE"
	LIMITS=$(cat << EOF
# $FILE for qemu-web-desktop
*    soft memlock 20000000
*    hard memlock 20000000
@kvm soft memlock unlimited
@kvm hard memlock unlimited
EOF
)
	# append only if token not found in FILE
	grep -q "kvm soft memlock unlimited" "$FILE" || echo "$LIMITS" >> "$FILE"

	# update apache2 mem allocation limit
	if [ -e $PREFIX/etc/systemd/system/httpd.service ]; then
		DIR=$PREFIX/etc/systemd/system/httpd.service.d/
	else
		DIR=$PREFIX/etc/systemd/system/apache2.service.d/
	fi
	mkdir -p $DIR
	FILE=$DIR/override.conf
	echo "Creating  $FILE"
	dd status=none of=${FILE} << EOF
# $FILE for qemu-web-desktop
[Service]
LimitMEMLOCK=infinity
EOF

	echo "Updating  GRUB boot, kernel modules, rules"
	if [ "x$DIST" = "xDebian" ]; then
		update-initramfs -u
		update-grub
	else
		mkinitcpio -p linux
		grub-mkconfig -o /boot/grub/grub.cfg || echo "WARNING: Failed to update GRUB"
	fi
	udevadm control --reload-rules
	udevadm trigger
	systemctl daemon-reload

	# request to uncomment GPU section in index.html
	echo "-----------------------------------------------------------------"
	echo "Check the above files, then UNCOMMENT the GPU section in the file"
	echo "  $index_html"
	echo "You may use 'sudo $p edit web' to do so."
	echo "To release the GPU back to the server use: 'sudo $p gpu_unlock'"
	echo "Now please reboot to activate changes."
	echo "-----------------------------------------------------------------"
	;;

# ------------------------------------------------------------------------------
	  edit|--edit)
	# default to edit machine file
	FILE=$machine_conf
	case "$2" in
		machine|machines|vm|list)
		;;
		landing|web|index|html)
		FILE=$index_html
		;;
		config|service)
		FILE=$config_pl
		;;
	esac
	echo "Edit $FILE"
	editor "$FILE" || if [ "x$EDITOR" = "x" ]; then edit "$FILE"; else $EDITOR "$FILE"; fi
	# trigger 'download' command
	$p download
	;;
	
# ------------------------------------------------------------------------------

	  version|--version|-v)
	$cgi -v
	exit 1
	;;

	  help|--help|-h)
	echo "usage: qwdctl [help|download|refresh|status|start VM|stop|gpu|edit] ..."
	echo " "
	echo "  controls the availability of virtual machines for the "
	echo "  qemu-web-desktop/DARTS service."
	echo "  In addition, the status of the running sessions can be displayed"
	echo "  and it is possible to start virtual machines manually and display"
	echo "  them in a browser."
	echo 
	echo "Syntax:"
	echo "   qwdctl --edit | --download | --start VM | --status"
	echo 
	echo "OPTIONS:"
	echo "  --download|download|update"
	echo "      scan the $machine_conf file for [name.ext] and download them when URL are given."
	echo "      a 'refresh' is then performed. Virtual machine images are stored into $qwdprefix/machines."
	echo " "
	echo "  --refresh|refresh"
	echo "      scan the $machine_conf file, and generate the $machine_html that lists"
	echo "      available images to show in the qemu-web-desktop main form."
	echo " "
	echo "  --status|status"
	echo "      list running sessions."
	echo " "
	echo "  --start VM|start VM ..."
	echo "      start the given VM or ISO file (full path) in a browser. Connect to it with"
	echo "      the displayed URL. Further arguments are passed to the service (see config.pl)"
	echo "      e.g. --snapshot_alloc_mem=1024 (in GB) and --snapshot_alloc_cpu=2."
	echo "      Changes are lost except when specifying option --snapshot_use_master=1,"
  echo "      requiring write access. When running ISO's you may also specify"
  echo "      --snapshot_alloc_disk=40 (in GB)."
  echo "      Use 'qwdctl --start' to get the full list of VM settings."
	echo " "
	echo "  --stop|stop TOKEN"
	echo "      stop sessions matching TOKEN. Some snapshot files may be left-over."
	echo " "
	echo "  --gpu|gpu VENDOR:MODEL"
	echo "      configure GPU for pass-through. The GPU_ids are e.g. '10de:1d01'"
	echo " "
	echo "  --gpu_unlock|gpu_unlock"
	echo "      unlock/re-attach all GPU's to server (uninstall pass-through)."
	echo " "
	echo "  --edit|edit [machines|config|web]"
	echo "      edit the VM/machine list, the service configuration file or the service web page."
	echo "      In the case of the VM list, the '$p download' command is triggered automatically after edit."
	echo "      Set the \$EDITOR variable to select the text editor to use."
	echo "      In case the '$p edit machines' has no effect on the index.html service landing page,"
	echo "      specify set 'machines_insert=yes' in $p."
	echo 
	echo "  --version|version|-v"
	echo "      show qemu-web-desktop version"
	echo 
	echo "  --help|help|-h"
	echo "      show this help"
	echo 
	echo "AUTHOR:"
	echo "  Written by Roland Mas and Emmanuel Farhi"
	echo 
	echo "FILES:"
	echo "  - /etc/qemu-web-desktop/config.pl"
	echo 
	echo "  - /var/lib/qemu-web-desktop/machines"
	echo 
	echo "  - /usr/share/qemu-web-desktop/html/desktop"
	echo " "
	echo "  - $index_html"
	echo 
	echo "  - $machine_conf."
	echo "  Entries should contain lines:"
	echo "    [name.ext]"
	echo "    description=<name of machine to appear in the form>"
	echo "  In addition, any line with:"
	echo "    url=<link>"
	echo "  will retrieve the given file upon: $p download"
	echo "  Supported virtual machine formats: ISO, QCOW2, VDI, VMDK, RAW, VHD/VHDX, QED" 
	echo 
	echo "  - /var/lib/qemu-web-desktop/machines.html"
	echo " "
	echo "  - https://gitlab.com/soleil-data-treatment/soleil-software-projects/qemu-web-desktop"
	echo " "
	echo "ENVIRONMENT VARIABLES:"
	echo "  EDITOR   Set the text editor to use, e.g. nano. Default is to use 'editor' or \$EDITOR or 'edit'."
	echo " "

	exit 1
	;;
	  *)
	  echo "Usage: $p [help|download|refresh|status|start|stop|gpu|edit] ..."
	  exit 1
esac

# ------------------------------------------------------------------------------

