#! /bin/sh

#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.
# 
# qwdctl download|update
#  scan the /etc/qemu-web-desktop/machines.conf file for [name.ext] and download them when URL are given.
#  a 'refresh' is then performed. Virtual machine images are stored into /var/lib/qemu-web-desktop/machines.
# 
# qwdctl refresh
#  scan the /etc/qemu-web-desktop/machines.conf file, and generate the 
#  /var/lib/qemu-web-desktop/include.html that lists
#  available images to show in the qemu-web-desktop main form.
#
# qwdctl status
#  list running sessions.
#
# qwdctl stop TOKEN
#  stop sessions matching TOKEN. Some snapshot files may be left-over.
#
# qwdctl gpu [VENDOR:MODEL]
#  configure GPU for pass-through. The GPU_ids are e.g. '10de:1d01' or '10de:1d01,10de:1d01'. 
#  When not given, the ID is requested.
#
# qwdctl gpu_unlock
#  re-attach/unlock all GPU's to server (uninstall pass-through).
#
# qwdctl edit [machine|config|web]
#  edit the list of machines (VM), the service configuration file or the service web page.
#  In the case of the VM list, the qwdctl download is triggered automatically after edit.

set -e

# file to process
machine_conf=/etc/qemu-web-desktop/machines.conf
index_html=/usr/share/qemu-web-desktop/html/desktop/index.html
config_pl=/etc/qemu-web-desktop/config.pl

# generated files, should be linked into /usr/share/qemu-web-desktop/html/desktop/
qwdprefix=/var/lib/qemu-web-desktop
machine_html=$qwdprefix/machines.html

case "$1" in
# ------------------------------------------------------------------------------
    download|update)
	mkdir -p $qwdprefix/machines || true
	cd $qwdprefix/machines

	for i in $(confget -f $machine_conf -q sections) ; do
	    mkdir -p downloads/$i || true
	    u=$(confget -f $machine_conf -s $i url)
	    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
	$0 refresh
	;;
# ------------------------------------------------------------------------------
    refresh)
	mkdir -p $qwdprefix/snapshots || true
	chown _qemu-web-desktop $qwdprefix/snapshots
	mkdir -p $qwdprefix/machines || true
	cd $qwdprefix/machines
	
	# list of machines
	t=$(mktemp $machine_html.XXXXXX)
	chmod 644 $t
	chmod a+rx $qwdprefix
	for i in $(confget -f $machine_conf -q sections) ; do
	    d=$(confget -f $machine_conf -s $i description)
	    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
	;;
# ------------------------------------------------------------------------------
	  status)
	echo "Active sessions:"
	echo "session_ID:user:VM\t| #cpu\t| #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 -Ts'\t')
	u=$(echo "$table" | uniq )
	echo "$u"
	;;
# ------------------------------------------------------------------------------
	  stop)
	if [ "x$2" = "x" ]; then
	  echo "Usage: $0 stop <token>";
	  exit 1;
	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)
	
	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/initramfs-tools/modules                        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"
	echo "  /etc/systemd/system/apache2.service.d/override.conf remove"
	read -p "Continue (y/N)?" choice
	case "$choice" in 
		y|Y ) echo "Proceeding";;
		*) echo "Aborting." && exit 1;;
	esac
	
	# restore default GRUB options
	FILE=/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
	FILE=/etc/initramfs-tools/modules
	echo "Updating  $FILE (old stored as $FILE.$DATE)"
	cp $FILE $FILE.$DATE
	sed -i "s/^vfio/#vfio/g"                 $FILE
	sed -i "s/^vhost-netdev/#vhost-netdev/g" $FILE

	rm -f /etc/modprobe.d/vfio.conf
	rm -f /etc/udev/rules.d/10-qemu-hw-users.rules
	rm -f /etc/systemd/system/apache2.service.d/override.conf
	
	update-initramfs -u
	update-grub
	udevadm control --reload-rules
	udevadm trigger
	systemctl daemon-reload
	
	# request to uncomment GPU section in index.html
	FILE=$index_html
	echo "-----------------------------------------------------------------"
	echo "Check the above files, then COMMENT the GPU section in the file"
	echo "  $FILE"
	echo "You may use 'sudo $0 edit web' to do so."
	echo "Then reboot..."
	echo "-----------------------------------------------------------------"
	;;
	
# ------------------------------------------------------------------------------
	  gpu|gpu_lock)
	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: $0 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:"
	echo "  /etc/default/grub      (GRUB_CMDLINE_LINUX_DEFAULT) update"
	echo "  /etc/initramfs-tools/modules                        append"
	echo "  /etc/security/limits.conf                           append"
	echo "  /etc/modprobe.d/vfio.conf                           create"
	echo "  /etc/udev/rules.d/10-qemu-hw-users.rules            create"
	echo "  /etc/systemd/system/apache2.service.d/override.conf create"
	read -p "Continue (y/N)?" choice
	case "$choice" in 
		y|Y ) echo "Proceeding";;
		*) echo "Aborting." && exit 1;;
	esac

	# GRUB
	FILE=/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: $0 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=/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
	FILE=/etc/initramfs-tools/modules
	MODULES=$(cat << EOF
# /etc/initramfs-tools/modules for qemu-web-desktop
vfio
vfio_iommu_type1
vfio_pci
vfio_virqfd
vhost-netdev
EOF
)
	# append only if token not found in FILE
	VFIO_there=no
	grep -q "^vfio_pci" "$FILE" && VFIO_there=yes
	if [ "x$VFIO_there" = "xyes" ]; then
		echo "WARNING: The VFIO modules seem to already be in place in $FILE. Keeping untouched."
	else
		echo "Appending $FILE (old stored as $FILE.$DATE)"
		cp $FILE $FILE.$DATE
		echo "$MODULES" >>"$FILE"
	fi

	# udev Rules
	FILE=/etc/udev/rules.d/10-qemu-hw-users.rules
	echo "Creating  $FILE"
	dd status=none of=${FILE} << EOF
# /etc/udev/rules.d/10-qemu-hw-users.rules
SUBSYSTEM=="vfio", OWNER="root", GROUP="kvm"
EOF
  
	# limits: append at end
	FILE=/etc/security/limits.conf
	echo "Appending $FILE"
	LIMITS=$(cat << EOF
# /etc/security/limits.conf
*    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
	mkdir -p /etc/systemd/system/apache2.service.d/
	FILE=/etc/systemd/system/apache2.service.d/override.conf
	echo "Creating  $FILE"
	dd status=none of=${FILE} << EOF
# /etc/systemd/system/apache2.service.d/override.conf
[Service]
LimitMEMLOCK=infinity
EOF

	echo "Updating  GRUB boot, kernel modules, rules"
	update-initramfs -u
	update-grub
	udevadm control --reload-rules
	udevadm trigger
	systemctl daemon-reload

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

# ------------------------------------------------------------------------------
	  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"
	# trigger 'download' command
	$0 download
	;;
	
# ------------------------------------------------------------------------------
	  *)
	echo "Unknown command $1"
	echo "Usage: $0 controls the availability of virtual machines for the qemu-web-desktop/DARTS service."
	echo "  The main file to tune is $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. "
	echo "  Supported virtual machine formats include: ISO, QCOW2, VDI, VMDK, RAW, VHD/VHDX, QED" 
	echo " "
	echo "$0 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 "$0 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 "$0 status"
	echo "  list running sessions"
	echo " "
	echo "$0 stop TOKEN"
	echo "  stop sessions matching TOKEN. Some snapshot files may be left-over."
	echo " "
	echo "$0 gpu VENDOR:MODEL"
	echo "  configure GPU for pass-through. The GPU_ids are e.g. '10de:1d01'"
	echo " "
	echo "$0 gpu_unlock"
	echo "  unlock/re-attach all GPU's to server (uninstall pass-through)."
	echo " "
	echo "$0 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 '$0 download' command is triggered automatically after edit."
	echo 
	exit 1
esac

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

