#!/bin/bash

#Qemu wrapper to handle LUKS & raw disk images and high performance (KVM, OpenGL, VirtIO) setup.
#By Frank Endres <frankendres@tuxfamily.org> - disabled, use first@last.fr
#Copyright 2023-04-30, under CeCILL / GPL licence

CRYPTSETUP="/sbin/cryptsetup"
SFDISK="/sbin/sfdisk" #FIXME: pb with GPT (doesn't work with sgdisk either)
LOSETUP="/sbin/losetup"
IP="/sbin/ip"


#============================== CHECK ==============================
function printUsage() {
	echo ""
	echo "This Qemu helper script features:"
	echo "- mounting of raw images partitions;"
	echo "with support of LUKS (Linux Unified Key Setup) disk images;"
	echo "- high performance (KVM, OpenGL, VirtIO) with tap / bridged networking setup."
	echo ""
	echo "usage: $0 create disk.luks size"
	echo "example: $0 create slackware.luks 10G"
	echo ""
	echo "usage: $0 config [disk.(raw|qcow2|luks|iso)]... [/share]"
	echo "example: $0 config slackware.iso slackware.luks > vm.qemu"
	echo "usage: $0 run vm.qemu"
	echo ""
	echo "usage: $0 mount disk.(raw|luks) part_num [mount_point]"
	echo "example: $0 mount slackware.luks 2 /mnt"
	echo "usage: $0 umount /dev/loopdevice"
	echo ""
}

action=$1
if [ "$action" != "create" ] && [ "$action" != "config" ] && [ "$action" != "run" ] && [ "$action" != "mount" ] && [ "$action" != "umount" ]; then
	printUsage
	exit 1
fi
if [ "$action" == "create" ] && (( $# != 3 )); then
	echo "error: bad parameters count"
	printUsage
	exit 1
fi
if [ "$action" == "run" ] && (( $# != 2 )); then
	echo "error: bad parameters count"
	printUsage
	exit 1
fi
if [ "$action" == "mount" ] && (( $# != 4 )) && (( $# != 3 )); then
	echo "error: bad parameters count"
	printUsage
	exit 1
fi
if [ "$action" == "umount" ] && (( $# != 2 )); then
	echo "error: bad parameters count"
	printUsage
	exit 1
fi


if [ "$action" == "create" ]; then
	img=$2
	size=$3
	if [ -f "$img" ]; then
		echo "error: file '$img' already exists; won't overwrite"
		exit 3
	fi
	fmt=$(echo "$img" | rev | cut -f1 -d. | rev)
	if [ "$fmt" != "raw" ] && [ "$fmt" != "qcow2" ] && [ "$fmt" != "luks" ]; then
		echo "error: unsupported '$fmt' disk image format"
		printUsage
		exit 2
	fi
fi

if [ "$action" == "config" ]; then
	shift
	smb=0
	images=""
	while [ ! -z "$1" ]; do
		img="$1"
		if [ -d "$img" ]; then
			if [ "$smb" == "1" ]; then
				echo "error: only one SMB share is supported"
				printUsage
				exit 2
			fi
			smb=1
		elif [ ! -f "$img" ]; then
			echo "error: file '$img' not found"
			printUsage
			exit 2
		fi
		images="$images $img"
		shift
	done
fi

if [ "$action" == "run" ]; then
	cfg=$2
	if [ ! -f "$cfg" ]; then
		echo "error: file '$cfg' not found"
		printUsage
		exit 2
	fi
fi

if [ "$action" == "mount" ]; then
	img=$2
	num=$3
	mnt=${4%/} #last parameter, with '/' suffix removed
	if [ ! -f "$img" ]; then
		echo "error: file '$img' not found"
		printUsage
		exit 2
	fi
	if [ ! -z "$mnt" ] && [ ! -d "$mnt" ]; then
		echo "error: '$mnt' is not a directory"
		printUsage
		exit 2
	fi
fi

if [ "$action" == "umount" ]; then
	dev=$2
	if [ ! -b $dev ] || ! echo $dev | grep -q "^/dev/loop"; then
		echo "error; '$dev' is not a loop device"
		printUsage
		exit 2
	fi
fi


#============================== DO ==============================

if [ "$action" == "create" ]; then
	if [ "$fmt" == "qcow2" ]; then
		qemu-img create -f qcow2 "$img" $size
	else
		qemu-img create -f raw "$img" $size
		if [ "$fmt" == "luks" ]; then
			sudo $CRYPTSETUP luksFormat "$img"
		fi
	fi


elif [ "$action" == "config" ]; then
	cat << EOF
qemu-system-x86_64
-accel kvm -cpu host -smp 2 -m 2048
-vga virtio -display gtk,gl=on,full-screen=on
#-vga virtio -display sdl,gl=on
-audio pa,model=hda #no audio when root
#-audio sdl,model=hda
EOF
	printf -v macaddr "52:54:%02x:%02x:%02x:%02x" $(( $RANDOM & 0xff)) $(( $RANDOM & 0xff )) $(( $RANDOM & 0xff)) $(( $RANDOM & 0xff ))
	if [ "$smb" == "0" ]; then
		echo "#-nic user,model=virtio-net-pci" #network user mode
		echo "-nic tap,model=virtio-net-pci,mac=$macaddr,script=no,downscript=no"
	fi
	if [ ! -f OVMF.fd ]; then echo -n "#"; fi
	echo "-drive if=pflash,format=raw,file=OVMF.fd #cp /usr/share/OVMF/x64/OVMF.fd ."
	for img in $images; do
		if [ -d "$img" ]; then
			echo "-nic user,model=virtio-net-pci,smb=$img"
			echo "#-nic tap,model=virtio-net-pci,mac=$macaddr,script=no,downscript=no"
		else
			fmt=$(echo "$img" | rev | cut -f1 -d. | rev)
			if [ "$fmt" == "iso" ]; then
				echo "-drive media=cdrom,file=$img -boot d"
			elif [ "$fmt" == "raw" ] || [ "$fmt" == "luks" ]; then #raw format should be explicit
				echo "-drive driver=raw,file=$img,if=virtio,aio=native,cache.direct=on"
			else #auto-detect format
				echo "-drive file=$img,if=virtio,aio=native,cache.direct=on"
			fi
		fi
	done


elif [ "$action" == "run" ]; then
	br=$($IP link show type bridge | head -n1 | cut -f2 -d:)
	if [ -z "$br" ] && cat "$cfg" | grep -q "^-nic tap"; then
		echo "error: tap networking enabled without available bridge"
		echo "to fix: enable user networking instead in '$cfg'"
		echo "or run: /usr/libexec/lxc/lxc-net start"
		exit 3
	fi
	
	for img in $(cat "$cfg" | sed -n -e 's/#.*$//' -e 's/.*file=\(.*\.luks\),.*/\1/' -e /\.luks/p); do
		name=$(basename $img | sed 's/\.luks$//')
		if [ ! -b "/dev/mapper/$name" ]; then
			sudo $CRYPTSETUP luksOpen "$img" "$name" || exit
			sudo chown $USER /dev/mapper/$name
		fi
	done
	
	if cat "$cfg" | grep -q "^-nic tap"; then
		lastnum=$($IP tuntap list | cut -d: -f1 | sort | cut -c4-)
		if [ -z "$lastnum" ]; then tapnum=0; else let tapnum=$lastnum+1; fi
		tap=tap$tapnum
		sudo $IP tuntap add mode tap $tap
		sudo $IP link set dev $tap master $br
		sudo $IP link set $tap up promisc on
		#~ sudo $IP addr flush dev $tap ?
	fi
	
	qemu=$(cat "$cfg" | sed -e 's/#.*$//' -e 's@file=\(.*/\)*\([^/]\+\)\.luks,@file=/dev/mapper/\2,@' -e s/^-nic\ tap/-nic\ tap,ifname=$tap/ | tr '\n' ' ')
	echo $qemu
	$qemu
	
	if cat "$cfg" | grep -q "^-nic tap"; then sudo $IP tuntap del $tap mode tap; fi
	for img in $(cat "$cfg" | sed -n -e 's/#.*$//' -e 's/.*file=\(.*\.luks\),.*/\1/' -e /\.luks/p); do
		name=$(basename $img | sed 's/\.luks$//')
		sudo $CRYPTSETUP luksClose "$name"
	done


elif [ "$action" == "mount" ]; then
	fmt=$(echo "$img" | rev | cut -f1 -d. | rev)
	if [ "$fmt" == "luks" ]; then
		name=$(basename $img | sed 's/\.luks$//')
		if [ ! -b "/dev/mapper/$name" ]; then sudo $CRYPTSETUP luksOpen "$img" "$name"; fi
		sudo chown $USER /dev/mapper/$name
		dm=$(readlink -f "/dev/mapper/$name")
		#~ sectorSize=$($SFDISK -l $dm | grep " = " | cut -f2 -d= | sed 's/[^0-9]//g')
		startSector=$($SFDISK -l $dm | grep "[a-z]$num " | awk '{ print $2 }')
		if [ -z "$startSector" ]; then
			echo "error; cant' find '$num' partition"
			sudo $CRYPTSETUP luksClose $dm
			exit 3
		fi
		#~ let offset=$sectorSize*$startSector
		let offset=512*$startSector
		sudo $LOSETUP -f -o $offset $dm
		dev=$($LOSETUP -a | grep "($dm)" |cut -f1 -d:)
	else
		#~ sectorSize=$($SFDISK -l "$img" | grep " = " | cut -f2 -d= | sed 's/[^0-9]//g')
		startSector=$($SFDISK -l "$img" | grep "[a-z]$num " | awk '{ print $2 }')
		#~ let offset=$sectorSize*$startSector
		if [ -z "$startSector" ]; then
			echo "error; cant' find '$num' partition"
			exit 3
		fi
		let offset=512*$startSector
		sudo $LOSETUP -f -o $offset "$img"
		dev=$($LOSETUP -a | grep "$img" |cut -f1 -d:)
	fi
	
	if [ ! -z "$mnt" ]; then
		sudo mount $dev $mnt
	fi
	echo $dev
	

elif [ "$action" == "umount" ]; then
	if $LOSETUP -a | grep "^$dev" | grep -q '(/dev/dm-'; then
		dm=$($LOSETUP -a | grep "^$dev" | cut -f2 -d'(' | cut -f1 -d')')
	fi
	if mount | grep -q "^$dev on"; then sudo umount $dev; fi
	if $LOSETUP -a | grep -q "^$dev"; then sudo $LOSETUP -d $dev; fi
	if [ ! -z "$dm" ]; then
		sudo $CRYPTSETUP luksClose $dm
	fi
fi