#!/bin/sh
#-------------------------------------------------------------------------+
# Copyright (C) 2015 Matt Churchyard (churchers@gmail.com)
# All rights reserved
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted providing that the following conditions
# are met:
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
# IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY
# DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
# OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
# STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING
# IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

# set us up for zfs use
# we need the zfs dataset name for zfs commands, and the file system path
# for bhyve. It's highly possible these will be different.
# We ask user to specify the dataset name as "zfs:pool/dataset"
# This can then be used for zfs commands, and we can retrieve the mountpoint
# to find out the file path for bhyve
#
# This then overwrites $vm_dir with the file system path, so that the 
# rest of vm-bhyve can work normally, regardless of whether we are in zfs mode
# or not
#
# If zfs is enabled, the following global variables are set
# VM_ZFS=1
# VM_ZFS_DATASET={pool/dataset}
#
# @modifies VM_ZFS VM_ZFS_DATASET vm_dir
#
__zfs_init(){
	local _zfs

	# check for zfs storage location
	# user should specify "zfs:pool/dataset" if they want ZFS support
	_zfs=$(echo "${vm_dir}" | cut -c -3)

	# are we storing on ZFS?
	if [ ${_zfs} = 'zfs' ]; then

		# check zfs running
		kldstat -n zfs >/dev/null 2>&1
		[ $? -ne 0 ] && __err "ZFS support requested but ZFS not available"

		# global zfs details
		VM_ZFS="1"
		VM_ZFS_DATASET=$(echo "${vm_dir}" | cut -c 5-)

		# update vm_dir
		# this makes sure it exists, confirms it's mounted & gets correct path in one go
		vm_dir=$(mount | grep "^${VM_ZFS_DATASET} " |cut -d' ' -f3)
		[ -z "${vm_dir}" ] && __err "unable to locate mountpoint for ZFS dataset ${VM_ZFS_DATASET}"
	fi
}

# make a new dataset
# this is always called when creating a new vm, but will do nothing
#
# @param string _name name of the dataset (under VM_ZFS_DATASET) to create
#
__zfs_make_dataset(){
	local _name="$1"
	local _opts="$2"

	if [ -n "${_name}" -a "${VM_ZFS}" = "1" ]; then
		__zfs_format_options "_opts" "${_opts}"
		zfs create ${_opts} "${VM_ZFS_DATASET}/${_name}"
		[ $? -ne 0 ] && __err "failed to create new ZFS dataset ${VM_ZFS_DATASET}/${_name}"
	fi
}

# destroy a dataset
#
# @param string _name name of the dataset to destroy (under VM_ZFS_DATASET)
#
__zfs_destroy_dataset(){
	local _name="$1"

	if [ -n "${_name}" -a "${VM_ZFS}" = "1" ]; then
		zfs destroy -rf "${VM_ZFS_DATASET}/${_name}" >/dev/null 2>&1
		[ $? -ne 0 ] && __err "failed to destroy ZFS dataset ${VM_ZFS_DATASET}/${_name}"
	fi
}

# rename a dataset
# as with other zfs functions, the arguments should just be the name
# of the dataset under $VM_ZFS_DATASET (in this case just guest name)
#
# @param string _old the name of the dataset to rename
# @param string _new the new name
#
__zfs_rename_dataset(){
	local _old="$1"
	local _new="$2"

	if [ -n "${_old}" -a -n "${_new}" -a "${VM_ZFS}" = "1" ]; then
		zfs rename "${VM_ZFS_DATASET}/${_old}" "${VM_ZFS_DATASET}/${_new}" >/dev/null 2>&1
		[ $? -ne 0 ] && __err "failed to rename ZFS dataset ${VM_ZFS_DATASET}/${_old}"
	fi
}

# make a zvol for a guest disk image
#
# @param string _name name of the guest (will be a dataset under VM_ZFS_DATASET)
# @param string _size how big to create the dataset
# @param int _sparse=0 set to 1 for a sparse zvol
#
__zfs_make_zvol(){
	local _name="$1"
	local _size="$2"
	local _sparse="$3"
	local _user_opts="$4"
	local _opt="-V"

	[ ! "${VM_ZFS}" = "1" ] && __err "cannot use ZVOL storage unless ZFS support is enabled"
	[ "${_sparse}" = "1" ] && _opt="-sV"

	__zfs_format_options "_user_opts" "${_user_opts}"
	zfs create ${_opt} ${_size} -o volmode=dev ${_user_opts} "${VM_ZFS_DATASET}/${_name}"
	[ $? -ne 0 ] && __err "failed to create new ZVOL ${VM_ZFS_DATASET}/${_name}"
}

# format options for zfs commands
# options are stored in configuration separated by |
# we need to replace that with -o
#
# @modifies $_val
#
__zfs_format_options(){
	local _val="$1"
	local _c_opts="$2"

	if [ -n "${_c_opts}" ]; then
		_c_opts=$(echo "${_c_opts}" |sed -e 's/\ / -o /')
		_c_opts="-o ${_c_opts}"
		setvar "${_val}" "${_c_opts}"
		return 0
	fi

	setvar "${_val}" ""
}

# 'vm snapshot'
# create a snapshot of a guest
# specify the snapshot name in zfs format guest@snap
# if no snapshot name is specified, Y-m-d-H:M:S will be used
# 
# @param flag (-f) force snapshot if guest is running
# @param string _name the name of the guest to snapshot
#
__zfs_snapshot(){
	local _name _snap _opt _force _snap_exists

	while getopts f _opt; do
		case $_opt in
			f)
				_force=1
				;;
			*)
				__usage
				;;
		esac
	done

	shift $((OPTIND - 1))
	_name="$1"

	# try to get snapshot name
	# we support normal zfs syntax for this
	_snap_exists=$(echo "${_name}" | grep "@")

	if [ -n "${_snap_exists}" ]; then
		_snap=${_name##*@}
		_name=${_name%%@*}
	fi

	[ -z "${_name}" ] && __usage
	[ ! "${VM_ZFS}" = "1" ] && __err "cannot snapshot guests unless ZFS support is enabled"
	[ ! -e "${vm_dir}/${_name}/${_name}.conf" ] && __err "${_name} does not appear to be an existing virtual machine"
	[ -z "${_snap}" ] && _snap=$(date +"%Y-%m-%d-%H:%M:%S")

	if ! __vm_confirm_stopped "${_name}" >/dev/null; then
		[ -z "${_force}" ] && __err "${_name} must be powered off first (use -f to override)"
	fi

	zfs snapshot -r ${VM_ZFS_DATASET}/${_name}@${_snap}
	[ $? -ne 0 ] && __err "failed to create recursive snapshot of virtual machine"
}

# 'vm rollback'
# roll a guest back to a previous snapshot
# we show zfs errors here as it will fail if the snapshot is not the most recent.
# zfs will output an error mentioning to use '-r', and listing the snapshots
# that will be deleted. makes sense to let user see this and just support 
# that option directly
#
# @param flag (-r) force deletion of more recent snapshots
# @param string _name name of the guest
#
__zfs_rollback(){
	local _name _snap _opt _force _fs _snap_exists

	while getopts r _opt; do
		case $_opt in
			r)
				_force="-r"
				;;
			*)
				__usage
				;;
		esac
	done

	shift $((OPTIND - 1))

	_snap_exists=$(echo "${1}" | grep "@")
	[ -z "${_snap_exists}" ] && __err "a snapshot name must be provided in guest@snapshot format"

	_name="${1%%@*}"
	_snap="${1##*@}"

	[ -z "${_name}" -o -z "${_snap}" ] && __usage
	[ ! "${VM_ZFS}" = "1" ] && __err "cannot rollback guests unless ZFS support is enabled"
	[ ! -e "${vm_dir}/${_name}/${_name}.conf" ] && __err "${_name} does not appear to be an existing virtual machine"

	__vm_confirm_stopped "${_name}" || exit 1

	# list all datasets and zvols under guest
	zfs list -o name -rHt filesystem,volume ${VM_ZFS_DATASET}/${_name} | \
	while read _fs; do
		zfs rollback ${_force} ${_fs}@${_snap}
		[ $? -ne 0 ] && exit $?
	done
}

# 'vm clone'
# clone a vm
# this makes a true zfs clone of the specifies guest
#
# @param string _old the guest to clone
# @param string _new name of the new guest
#
__zfs_clone(){
	local _old="$1"
	local _new="$2"
	local _uuid _uuidpart _fs _newfs

	[ -z "${_old}" -o -z "${_new}" ] && __usage
	[ ! "${VM_ZFS}" = "1" ] && __err "cannot clone guests unless ZFS support is enabled"
	[ ! -e "${vm_dir}/${_old}/${_old}.conf" ] && __err "${_old} does not appear to be an existing virtual machine"
	[ -e "${vm_dir}/${_old}/run.lock" ] && __err "${_old} must be powered off first"
	[ -d "${vm_dir}/${_new}" ] && __err "directory ${vm_dir}/${_new} already exists"

	_uuid=$(uuidgen)
	_uuidpart=$(echo "${_uuid}" |awk -F- '{print $1}')

	zfs snapshot -r ${VM_ZFS_DATASET}/${_old}@${_uuidpart}
	[ $? -ne 0 ] && __err "failed to create snapshot ${VM_ZFS_DATASET}/${_old}@${_uuidpart}"

	zfs list -o name -rHt filesystem,volume ${VM_ZFS_DATASET}/${_old} | \
	while read _fs; do
		_newfs=$(echo "${_fs}" | sed "s@${VM_ZFS_DATASET}/${_old}@${VM_ZFS_DATASET}/${_new}@")

		zfs clone ${_fs}@${_uuidpart} ${_newfs}
		[ $? -ne 0 ] && __err "error while cloning datasets"
	done

	# update new guest
	mv "${vm_dir}/${_new}/${_old}.conf" "${vm_dir}/${_new}/${_new}.conf"
	sysrc -inqf "${vm_dir}/${_new}/${_new}.conf" "uuid=${_uuid}" >/dev/null 2>&1
	rm "${vm_dir}/${_new}/vm-bhyve.log" >/dev/null 2>&1
}

# 'vm image create'
# create an image of a vm
# this creates an archive of the specified guest, stored in $vm_dir/images
# we use a uuid just in case we want to provide the ability to share images at any point
#
# @param optional string (-d) description of the image
# @param string _name name of guest to take image of
#
__zfs_image_create(){
	local _name
	local _opt _desc _uuid _date

	while getopts d: _opt ; do
		case $_opt in
			d)
				_desc=${OPTARG}
				;;
			*)
				__usage
				;;
		esac
	done

	shift $((OPTIND - 1))
	_name=$1
	_uuid=$(uuidgen)
	_date=$(date)

	[ -z "${_desc}" ] && _desc="No description provided"
	[ ! -e "${vm_dir}/${_name}" ] && __err "${_name} does not appear to be a valid virtual machine"

	# create the image dataset if we don't have it
	if [ ! -e "${vm_dir}/images" ]; then
		zfs create "${VM_ZFS_DATASET}/images" >/dev/null 2>&1
		[ $? -ne 0 ] && __err "failed to create image store ${VM_ZFS_DATASET}/images"
	fi

	# try to snapshot
	zfs snapshot -r "${VM_ZFS_DATASET}/${_name}@${_uuid}" >/dev/null 2>&1
	[ $? -ne 0 ] && __err "failed to create snapshot of source dataset ${VM_ZFS_DATASET}/${_name}@${_uuid}"

	# copy source
	echo "Creating a compressed image, this may take some time..."
	zfs send -R "${VM_ZFS_DATASET}/${_name}@${_uuid}" | xz > "${vm_dir}/images/${_uuid}.zfs.xz"
	[ $? -ne 0 ] && exit 1

	# done with the source snapshot
	zfs destroy ${VM_ZFS_DATASET}/${_name}@${_uuid}

	# create a description file
	sysrc -inqf "${vm_dir}/images/${_uuid}.manifest" "description=${_desc}" >/dev/null 2>&1
	sysrc -inqf "${vm_dir}/images/${_uuid}.manifest" "created=${_date}" >/dev/null 2>&1
	sysrc -inqf "${vm_dir}/images/${_uuid}.manifest" "name=${_name}" >/dev/null 2>&1
	sysrc -inqf "${vm_dir}/images/${_uuid}.manifest" "filename=${_uuid}.zfs.xz" >/dev/null 2>&1

	echo "Image of ${_name} created with UUID ${_uuid}"
}

# 'vm image provision'
# create a new vm from an image
#
# @param string _uuid the uuid of the image to use
# @param string _name name of the new guest
#
__zfs_image_provision(){
	local _uuid="$1"
	local _name="$2"
	local _file
	local _oldname

	[ -z "${_uuid}" -o -z "${_name}" ] && __usage
	[ ! -e "${vm_dir}/images/${_uuid}.manifest" ] && __err "unable to locate image with uuid ${_uuid}"
	[ -e "${vm_dir}/${_name}" ] && __err "directory ${vm_dir}/${_name} already exists"

	# get the data filename
	_file=$(sysrc -inqf "${vm_dir}/images/${_uuid}.manifest" filename)
	_oldname=$(sysrc -inqf "${vm_dir}/images/${_uuid}.manifest" name)
	[ -z "${_file}" -o -z "${_oldname}" ] && __err "unable to locate required details from the specified image manifest"
	[ ! -e "${vm_dir}/images/${_file}" ] && __err "image data file does not exist: ${vm_dir}/images/${_file}"

	# try to recieve
	echo "Unpacking compressed image, this may take some time..."
	cat "${vm_dir}/images/${_file}" | xz -d | zfs recv "${VM_ZFS_DATASET}/${_name}"
	[ $? -ne 0 ] && __err "errors occured while trying to unpackage the image file"

	# remove the original snapshot
	zfs destroy "${VM_ZFS_DATASET}/${_name}@${_uuid}" >/dev/null 2>&1

	# rename the guest configuration file
	mv "${vm_dir}/${_name}/${_oldname}.conf" "${vm_dir}/${_name}/${_name}.conf" >/dev/null 2>&1
	[ $? -ne 0 ] && __err "unpackaged image but unable to update guest configuration file"
}

# 'vm image list'
# list available images
#
__zfs_image_list(){
	local _file _uuid _ext
	local _format="%-38s %-16s %-30s %s\n"

	printf "${_format}" "UUID" "NAME" "CREATED" "DESCRIPTION"

	[ ! -e "${vm_dir}/images" ] && exit

	ls -1 ${vm_dir}/images/ | \
	while read _file; do
		_ext=$(echo "${_file}" | cut -d. -f2)

		if [ "${_ext}" = "manifest" ]; then
			_uuid=$(echo "${_file}" | cut -c -36)
			_desc=$(sysrc -inqf "${vm_dir}/images/${_uuid}.manifest" description)
			_created=$(sysrc -inqf "${vm_dir}/images/${_uuid}.manifest" created)
			_name=$(sysrc -inqf "${vm_dir}/images/${_uuid}.manifest" name)

			printf "${_format}" "${_uuid}" "${_name}" "${_created}" "${_desc}"
		fi
	done
}

# 'vm image destroy'
# destroy an image
#
# @param string _uuid the uuid of the image
#
__zfs_image_destroy(){
	local _uuid="$1"
	local _file

	[ -z "${_uuid}" ] && __usage
	[ ! -e "${vm_dir}/images/${_uuid}.manifest" ] && __err "unable to locate image with uuid ${_uuid}"

	# get the image filename
	_file=$(sysrc -inqf "${vm_dir}/images/${_uuid}.manifest" filename)
	[ -z "${_file}" ] && __err "unable to locate filename for the specified image"

	rm "${vm_dir}/images/${_uuid}.manifest"
	rm "${vm_dir}/images/${_file}"
}

# cmd 'vm image ...'
# parse the image command set
# these all rely on ZFS snapshots, so kept with zfs functions
#
# @param string _cmd the command after 'vm image '
#
__zfs_parse_image_cmd(){
	local _cmd="$1"
	shift

	# we only support these commands on zfs
	[ ! "${VM_ZFS}" = "1" ] && __err "the image command set is only available with ZFS storage"

	case "${_cmd}" in
		list)
			__zfs_image_list
			;;
		create)
			__zfs_image_create "$@"
			;;
		provision)
			__zfs_image_provision "$@"
			;;
		destroy)
			__zfs_image_destroy "$@"
			;;
		*)
			__usage
			;;
	esac
}
