#!/bin/bash
# sys/tool/vbox - VirtualBox management library

set -e

lib_load 'sys/tool/qemu-nbd'


######
# Native dependencies

sys_tool_vbox_client_packages=( rdesktop )
sys_tool_vbox_server_packages=( virtualbox virtualbox-ext-pack )


######
# Settings

sys_tool_vbox_config_init() {
	lib_setting_vars \
			vbox_mounttab \
			vbox_vm_dir \
			vbox_hd_dir vbox_hd_ext \
			vbox_iso_dir vbox_iso_ext \
			vbox_os_type \
			vbox_sata_controller vbox_ide_controller \
			vbox_mem_size vbox_vram_size \
			vbox_hd_size vbox_hd_port vbox_hd_device \
			vbox_dvd_port vbox_dvd_device

	lib_setting_arrays -ro vbox_data_types

	# VDI mounttab:
	#	$vm_name:$nbd_device:$nbd_partition:$mount_point0
	vbox_mounttab="$confdir/mounttab"

	# Location for storing VM data
	vbox_vm_dir="$datadir/machines"
	# Location for storing HD images
	vbox_hd_dir="$datadir/disks"
	vbox_hd_ext="vdi"
	# Location for storing ISO images
	vbox_iso_dir="$datadir/images"
	vbox_iso_ext="iso"

	# Default OS type (see vbox_manage list ostypes)
	vbox_os_type="Ubuntu_64"

	# Default SATA controller settings
	vbox_sata_controller="SATA Controller"
	# Default IDE controller settings
	vbox_ide_controller="IDE Controller"

	vbox_mem_size=512
	vbox_vram_size=64

	# Default size of new disk (in MB)
	vbox_hd_size=8096
	vbox_hd_port=0
	vbox_hd_device=0

	vbox_dvd_port=0
	vbox_dvd_device=0

	# vbox_data_types - List of registered data types
	vbox_data_types=( vm hd iso )
}

sys_tool_vbox_config_check() {
	vbox_data_init
}


######
# VirtualBox helpers

vbox_manage() { run VBoxManage "$@"; }
vbox_setproperty() { vbox_manage setproperty "$@"; }
vbox_snapshot() { vbox_manage snapshot "$@"; }

vbox_headless() { run VBoxHeadless "$@"; }


######
# Generic data storage helpers

# vbox_data_init - Initialize data directories
vbox_data_init() {
	local -a ddirs
	for dt in "${vbox_data_types[@]}"; do
		 eval ddirs=( "\${ddirs[@]}" "\$vbox_${dt}_dir" )
	done
	run_mkdir "${ddirs[@]}"
}

# vbox_data_list - List contents of directory stored in $1
vbox_data_list() {
	local var=$1
	(eval cd "\$$var" && ls -1)
}

# vbox_data_path - Construct a filename using a directory stored in $1 and
#   the name from $3 passed through a function named in $2
# $1 - data type
# $2 - name
vbox_data_filename() {
	has_args 2 "$@"
	local class=$1
	local var="vbox_${1}_dir"
	local func="vbox_${1}_name"
	local name
	name=$($func "$2")
	eval echo "\$$var/\$name"
}

vbox_data_filenames() {
	min_args 3 "$@"
	local var=$1
	local dt=$2
	local term=$3
	shift 3
	local name
	local -a _files
	for name in "$@"; do
		local file
		file=$(eval vbox_${dt}_filename "$name")
		_files=( "${_files[@]}" "$file" )
	done
	eval "$var=( "\${_files[@]}" )"
}
# vbox_data_exists - Returns success if data file exists
# $1 - data type
# $2 - name
vbox_data_exists() {
	has_args 2 "$@"
	local file
	file=$(vbox_data_filename "$@")
	test -f "$file"
}


# $1 - command line
# $2 - command verb
# $3 - data type (e.g. "hd")
# $4 - data term (e.g "image")
# ... - data names (to be transformed to filenames)
vbox_data_run() {
	local cmd=$1
	local verb=$2
	local dt=$3
	local term=$4
	shift 4
	local -a files
	vbox_data_filenames files "$dt" "$term" "$@"
	app_echo "$dt: $verb $term: ${files[@]}"
	run $cmd "${files[@]}"
}

# $1 - data type (e.g. "hd")
# $2 - data term (e.g "image")
# ... - names
vbox_data_rm() {
	min_args 3 "$@"
	vbox_data_run rm 'deleting' "$@"
}

# $1 - data type (e.g. "hd")
# $2 - data term (e.g "image")
# $3 - src
# $4 - tgt
vbox_data_mv() {
	has_args 4 "$@"
	if vbox_data_exists "$1" "$4"; then
		error "$1: $4: $2 exists"
	fi
	vbox_data_run mv 'renaming' "$@"
}

######
# HDs

vbox_hd_name() { echo "$1.vdi"; }
vbox_hd_filename() { vbox_data_filename hd "$1"; }
vbox_hd_exists() { vbox_data_exists hd "$1"; }

vbox_hd_create() {
	local name=$1
	local size=$2
	app_echo "hd: $name: creating ${size}MB image"
	local file
	file=$(vbox_hd_filename "$name")
	vbox_manage createhd --filename "$file" --size $size
}

vbox_hd_rename() { has_args 2 "$@"; vbox_data_mv hd 'disk image' "$1" "$2"; }
vbox_hd_delete() {
	min_args 1 "$@"
	local name=$1
	local delete=${2:-false}
	local file
	file=$(vbox_hd_filename "$name")
	local opt=
	! $delete || opt=--delete
	vbox_manage closemedium "$file" $opt
#	vbox_data_rm hd 'disk image' "$@"
}


######
# HD mounttab support

# vbox_mounttab_add - Add an entry to the mounttab
vbox_mounttab_add() {
	has_args 4 "$@"
	echo "$1:$2:$3:$4" >>"$vbox_mounttab"
}
# vbox_mounttab_remove - Remove an entry from the mounttab
vbox_mounttab_remove() {
	has_args 3 "$@"
	[ -f "$vbox_mounttab" ] || error "$vbox_mounttab does not exist"
	local new="$vbox_mounttab+"
	grep -v "^$1:$2:$3:" "$vbox_mounttab" >"$new" || true

	local old="$vbox_mounttab~"
	run mv "$vbox_mounttab" "$old"
	run mv "$new" "$vbox_mounttab"
}

# vbox_mounttab_devices - List of NBD devices being used
vbox_mounttab_devices() { cut -f2 -d ':' "$vbox_mounttab" | sort -nu; }
# vbox_mounttab_device - Find NBD device for given HD image
vbox_mounttab_device() {
	has_args 1 "$@"
	[ ! -f "$vbox_mounttab" ] && return
	grep "^$1:" "$vbox_mounttab" | head -n 1 | cut -f2 -d':'
}

# vbox_mounttab_device_next - Find next available NBD device
#   Starting after the last id, search through the device list
#   for the first unused device number.
vbox_mounttab_device_next() {
	max_args 1 "$@"
	local last=${1:--1}

	if [ ! -f "$vbox_mounttab" ]; then
		echo 0
		return
	fi

	local -a devs
	devs=( $(vbox_mounttab_devices) )
	# devse array index
	local dev=0
	# next device id
	local next=$(($last + 1))
	# skip devices that are bigger than the next id being considered
	while [ "${devs[$dev]}" ] && [ "$next" -gt "${devs[$dev]}" ]; do
		int_inc dev
	done
	# skip ids if the devices are in-use
	while [ "${devs[$dev]}" ] && [ "$next" -eq "${devs[$dev]}" ]; do
		int_inc dev next
	done

	if [ $next -eq $qemu_nbd_device_count ]; then
		error "all NBD devices are in use"
	fi

	echo $next
}

# vbox_mounttab_partitions - List all partition:path pairs for the image
vbox_mounttab_partitions() {
	has_args 1 "$@"
	[ ! -f "$vbox_mounttab" ] && return
	grep "^$1:" "$vbox_mounttab" | cut -f'3-4' -d':' | sort -n
}

# vbox_mounttab_partition - Find partition mount point
vbox_mounttab_partition() {
	has_args 3 "$@"
	[ ! -f "$vbox_mounttab" ] && return
	grep "^$1:$2:$3:" "$vbox_mounttab" | cut -f4 -d':'
}


######
# HD image mounting support

vbox_hd_image_partitions() {
	has_args 1 "$@"
	local name=$1

	local dev
	dev=$(vbox_mounttab_device "$name")

	# if we need to mount the image, remember we did it
	local mounted=true
	[ "$dev" ] && mounted=false || dev=$(vbox_hd_image_mount "$name")
	[ "$dev" ] || error "$name: unable to list partitions"

	# run fdisk on image device
	local file
	file=$(qemu_nbd_image_dev "$dev")
	run_sudo fdisk -l "$file"

	# if we mounted it, umount it
	$mounted && vbox_hd_image_umount "$dev"
}

# vbox_hd_mount_image - Internal helper for mounting HD image devices
vbox_hd_image_mount() {
	has_args 1 "$@"
	local name=$1

	local dev
	dev=$(vbox_mounttab_device "$name")
	if [ "$dev" ]; then
		warn "$name: image already mounted on nbd$dev"
		echo "$dev"
		return
	fi

	dev=$(vbox_mounttab_device_next)
	[ "$dev" ] || error "$name: no NBD devices available"

	local file
	file=$(vbox_hd_filename "$name")
	qemu_nbd_image_mount "$dev" "$file"

	echo "$dev"
}

# vbox_hd_image_umount - Internal helper for umounting an HD image device
vbox_hd_image_umount() { qemu_nbd_image_umount "$1"; }

# vbox_hd_partition_mount - Mount an image partition
vbox_hd_partition_mount() {
	min_args 2 "$@"
	local name=$1
	local part=$2

	local dev=$(vbox_mounttab_device "$name")
	if [ -z "$dev" ]; then
		info "$name: mounting image"
		dev=$(vbox_hd_image_mount "$name")
	fi

	if [ "$part" = "raw" ]; then
		vbox_mounttab_add "$name" "$dev" "$part" "<none>"
		echo $(qemu_nbd_image_dev "$dev")
		return
	fi

	has_args 3 "$@"
	local mnt=$3
	[ -d "$mnt" ] || run_mkdir "$mnt"
	mnt=$(realpath "$mnt")

	local mmnt
	mmnt=$(vbox_mounttab_partition "$name" "$dev" "$part")
	if [ "$mmnt" ]; then
		local level
		[ "$mnt" = "$mmnt" ] && level=warn || level=error
		$level "$name: partition $part already mounted at '$mmnt'"
		return
	fi
	if qemu_nbd_partition_mount "$dev" "$part" "$mnt"; then
		vbox_mounttab_add "$name" "$dev" "$part" "$mnt"
		app_echo "$name: partition $part mounted on $mnt"
	else
		error "$name: partition $part failed to mount on $mnt"
	fi
}

# vbox_hd_partition_umount - Unmount an image partition, unmounting
#   the image if no other partitions are using it.
vbox_hd_partition_umount() {
	has_args 2 "$@"
	local name=$1
	local part=$2

	local dev
	dev=$(vbox_mounttab_device "$name")
	if [ -z "$dev" ]; then
		warn "$name: image is not mounted"
		return
	fi

	if [ "$part" != "raw" ]; then
		local mnt
		mnt=$(vbox_mounttab_partition "$name" "$dev" "$part")
		if [ "$mnt" ]; then
			app_echo "$name: unmounting partition $part from $mnt"
			qemu_nbd_partition_umount "$mnt"
		else
			warn "$name: partition $part not mounted"
		fi
	fi

	vbox_mounttab_remove "$name" "$dev" "$part"

	# check the device number again
	local cdev
	cdev=$(vbox_mounttab_device "$name")
	if [ -z "$cdev" ]; then
		# no partitions remain mounted, unmount the image
		info "$name: no partitions in use, unmounting image"
		vbox_hd_image_umount "$dev"
	fi
}


######
# ISOs

vbox_iso_name() { echo "$1.iso"; }
vbox_iso_filename() { vbox_data_filename iso "$1"; }
vbox_iso_exists() { vbox_data_exists iso "$1"; }

vbox_iso_add() {
	min_args 2 "$@"
	local name=$1
	local iso=$2
	app_echo "iso: $name: uploading $iso"

	local file
	file=$(vbox_iso_filename "$name")
	run cp -a "$iso" "$file"
}

vbox_iso_rename() { has_args 2 "$@"; vbox_data_mv iso 'ISO image' "$1" "$2"; }
vbox_iso_delete() { min_args 1 "$@"; vbox_data_rm iso 'ISO image' "$@"; }


######
# VMs

vbox_vm_name() { echo "$1/$1.vbox"; }
vbox_vm_filename() { vbox_data_filename vm "$1"; }
vbox_vm_exists() { vbox_data_exists vm "$1"; }

vbox_vm_create() {
	min_args 1 "$@"
	local name=$1
	local register=${2:-true}

	local opt
	$register && opt="--register"

	info "$name: creating virtual machine"
	vbox_manage createvm --name "$name" --ostype $vbox_os_type $opt \
		--basefolder "$vbox_vm_dir"
}

vbox_vm_modify() {
	vbox_manage modifyvm "$@"
}

vbox_vm_register() {
	has_args 1 "$@"
	local name=$1
	vbox_manage registervm "$vbox_vmdir/$name"
}

vbox_vm_unregister() {
	has_args 2 "$@"
	local name=$1
	local delete=$2
	local opt=""
	local act="unregistered"
	if $delete; then
		opt="--delete"
		act="deleted"
	fi
	vbox_manage unregistervm "$name" $opt
	app_echo "$name: $act"
}

vbox_vm_config_rdp() {
	has_args 2 "$@"
	local name=$1
	local enabled=$2
	local value=on
	$enabled || value=off
	info "$name: setting RDP to $value"
	vbox_vm_modify "$name" --vrde $value
}


######
# VM State

vbox_vm_state_save() { has_args 1 "$@"; vbox_manage controlvm "$1" savestate; }
vbox_vm_state_adopt() { has_args 2 "$@"; vbox_manage adoptstate "$@"; }
vbox_vm_state_discard() { has_args 1 "$@"; vbox_manage discardstate "$@"; }


######
# VM Snapshots

vbox_vm_snapshot_list() { has_args 1 "$@"; vbox_snapshot "$1" list; }
vbox_vm_snapshot_info() { has_args 2 "$@"; vbox_snapshot "$1" showvminfo "$2"; }
vbox_vm_snapshot_take() { has_args 2 "$@"; vbox_snapshot "$1" take "$2"; }
vbox_vm_snapshot_delete() { has_args 2 "$@"; vbox_snapshot "$1" delete "$2"; }
vbox_vm_snapshot_restore() {
	has_args 2 "$@"
	if [ "$2" = current ]; then
		vbox_snapshot "$1" restorecurrent
	else
		vbox_snapshot "$1" restore "$2"
	fi
}


######
# SATA

vbox_sata_add() {
	has_args 1 "$@"
	local vm=$1
	app_echo "vm: $vm: sata: $vbox_sata_controller: creating controller"
	vbox_manage storagectl "$vm" --name "$vbox_sata_controller" \
		--add sata --controller IntelAHCI
}

vbox_hd_setup() {
	has_args 2 "$@"
	local vm=$1
	local hd=$2
	local file
	file=$(vbox_hd_filename "$hd")
	app_echo "vm: $vm: sata: $vbox_sata_controller: hd: media change: $hd"
	vbox_manage storageattach "$vm" \
		--storagectl "$vbox_sata_controller" \
		--port "$vbox_hd_port" --device "$vbox_hd_device" \
		--type hdd --medium "$file"
}


######
# IDE

vbox_ide_add() {
	has_args 1 "$@"
	local vm=$1
	app_echo "vm: $vm: ide: $vbox_ide_controller: creating controller"
	vbox_manage storagectl "$vm" --name "$vbox_ide_controller" \
		--add ide
}

vbox_dvd_media() {
	has_args 2 "$@"
	local vm=$1
	local iso=$2
	local file
	case "$iso" in
	({*}) file="$iso" ;;
	(eject) file=emptydrive ;;
	(none|emptydrive) file=$iso ;;
	(*) file=$(vbox_iso_filename "$iso") ;;
	esac

	app_echo "vm: $vm: ide: $vbox_ide_controller: dvd: media change: $iso"
	vbox_manage storageattach "$vm" \
		--storagectl "$vbox_ide_controller" \
		--port "$vbox_dvd_port" --device "$vbox_dvd_device" \
		--type dvddrive --medium "$file"
}