#!/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.

# 'vm list'
# list virtual machines
#
__vm_list(){
    local _name _guest _loader _cpu _memory _run _vm _auto _num
    local _format="%-15s %-10s %-11s %-6s %-9s %-12s %s\n"

    __vm_running_load
    printf "${_format}" "NAME" "GUEST" "LOADER" "CPU" "MEMORY" "AUTOSTART" "STATE"

    ls -1 "${vm_dir}" | \
    while read _name; do
        [ ! -e "${vm_dir}/${_name}/${_name}.conf" ] && continue

        __config_load "${vm_dir}/${_name}/${_name}.conf"
        __config_get "_guest" "guest"
        __config_get "_loader" "loader"
        __config_get "_cpu" "cpu"
        __config_get "_memory" "memory"

        __vm_running_check "_run" "${_name}"
        _num=1
        _auto="No"

        # find out if we auto-start this vm, and get sequence number
        for _vm in ${vm_list}; do
            [ "${_vm}" = "${_name}" ] && _auto="Yes [${_num}]"
            _num=$(($_num + 1))
        done

        # if stopped, see if it's locked by another host
        if [ "${_run}" = "Stopped" -a -e "${vm_dir}/${_name}/run.lock" ]; then
            _run=$(head -n1 "${vm_dir}/${_name}/run.lock")
            _run="Locked (${_run})"
        fi

        printf "${_format}" "${_name}" "${_guest}" "${_loader:-default}" "${_cpu}" "${_memory}" "${_auto}" "${_run}"
    done
}

# 'vm check name'
# check name of virtual machine
#
# @return int 0 if name is valid
#
__vm_check_name(){
    echo "$1" | egrep -iqs '^[a-z0-9][.a-z0-9-]{0,14}[a-z0-9]$'
}

# 'vm create'
# create a new virtual machine
#
# @param optional string (-t) _template the template to use (default = default)
# @param optional string (-s) _size guest size (default = 20G)
# @param string _name the name of the guest to create
#
__vm_create(){
    local _name _template _opt _size _vmdir _disk _disk_dev _uuid
    local _zfs_opts _num=0 _disk_size

    while getopts t:s: _opt ; do
        case $_opt in
            t) _template=${OPTARG} ;;
            s) _size=${OPTARG} ;;
            *) __usage ;;
        esac
    done

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

    [ -z "${_name}" ] && __usage
    __vm_check_name "${_name}" || __err "invalid virtual machine name - '${_name}'"

    # use default.conf template if none specified
    : ${_template:=default}

    [ ! -f "${vm_dir}/.templates/${_template}.conf" ] && \
        __err "unable to find template ${vm_dir}/.templates/${_template}.conf"

    # we need to get disk0 name and device type from the template
    __config_load "${vm_dir}/.templates/${_template}.conf"
    __config_get "_disk" "disk0_name"
    __config_get "_disk_dev" "disk0_dev"
    __config_get "_disk_size" "disk0_size" "20G"
    __config_get "_zfs_opts" "zfs_dataset_opts"

    # make sure template has a disk before we start creating anything
    [ -z "${_disk}" ] && __err "template is missing disk0_name specification"

    # get the directory to store our vm data in
    _vmdir="${vm_dir}/${_name}"
    [ -e "${_vmdir}" ] && __err "directory ${_vmdir} already exists"

    # if we're on zfs, make a new filesystem
    __zfs_make_dataset "${_name}" "${_zfs_opts}"

    [ ! -d "${_vmdir}" ] && mkdir "${_vmdir}" >/dev/null 2>&1
    [ ! -d "${_vmdir}" ] && __err "unable to create virtual machine directory ${_vmdir}"

    cp "${vm_dir}/.templates/${_template}.conf" "${_vmdir}/${_name}.conf"
    [ $? -ne 0 ] && __err "unable to copy template to virtual machine directory"

    # generate a uuid
    # jump through some hoops to make sure it gets its own line (but no gap)
    _uuid=$(uuidgen)
    __vm_config_set "${_name}" "uuid" "${_uuid}"

    # get any zvol options
    __config_get "_zfs_opts" "zfs_zvol_opts"

    # use cmd line size for disk 0 if specified
    [ -n "${_size}" ] && _disk_size="${_size}"

    # create each disk
    while [ -n "${_disk}" ]; do
        case "${_disk_dev}" in 
            zvol)
                __zfs_make_zvol "${_name}/${_disk}" "${_disk_size}" "0" "${_zfs_opts}"
                ;;			
            sparse-zvol)
                __zfs_make_zvol "${_name}/${_disk}" "${_disk_size}" "1" "${_zfs_opts}"
                ;;
            *)
                truncate -s "${_disk_size}" "${_vmdir}/${_disk}"
                [ $? -ne 0 ] && __err "failed to create sparse file for disk image"
                ;;
        esac

        # scrap size option from guest template
        sysrc -inxqf "${_vmdir}/${_name}.conf" "disk${_num}_size"

        # look for another disk
        _num=$((_num + 1))
        __config_get "_disk" "disk${_num}_name"
        __config_get "_disk_dev" "disk${_num}_dev"
        __config_get "_disk_size" "disk${_num}_size" "20G"
    done
}

# 'vm add'
# add a device to an existing guest
#
# @param string (-d) _device=network|disk the type of device to add
# @param string (-t) _type for disk, the type of disk - file|zvol|sparse-zvol
# @param string (-s) _sopt for disk the size, for network the virtual switch name
# @param string _name name of the guest
#
__vm_add(){
    local _name _device _type _sopt _opt

    while getopts d:t:s: _opt; do
        case $_opt in
            d)
                _device=${OPTARG}
                ;;
            t)
                _type=${OPTARG}
                ;;
            s)
                _sopt=${OPTARG}
                ;;
            *)
                __usage
                ;;
        esac
    done

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

    # check guest
    [ -z "${_name}" ] && __usage
    [ ! -e "${vm_dir}/${_name}/${_name}.conf" ] && __err "${_name} does not appear to be a valid virtual machine"

    case "${_device}" in
        disk)
            __vm_add_disk "${_name}" "${_type}" "${_sopt}"
            ;;
        network)
            __vm_add_network "${_name}" "${_sopt}"
            ;;
        *)
            __err "device must be one of the following: disk network"
            ;;
    esac
}

# add a disk to guest
# this creates the disk image or zvol and updates configuration file
# we use the same emulation as the existing disk(s)
#
# @private
# @param string _name name of the guest
# @param string _device type of device file|zvol|sparse-zvol
# @param string _size size of the disk to create
#
__vm_add_disk(){
    local _name="$1"
    local _device="$2"
    local _size="$3"
    local _num=0 _curr _diskname _emulation _zfs_opts

    : ${_device:=file}

    [ -z "${_size}" ] && __usage

    # get the last existing disk
    __config_load "${vm_dir}/${_name}/${_name}.conf"
    __config_get "_zfs_opts" "zfs_zvol_opts"

    while [ 1 ]; do
        __config_get "_curr" "disk${_num}_name"
        [ -z "${_curr}" ] && break
        __config_get "_emulation" "disk${_num}_type"
        _num=$((_num + 1))
    done

    [ -z "${_emulation}" ] && __err "failed to get emulation type of the existing guest disks"

    # create the disk first, then update config if no problems
    case "${_device}" in
        zvol)
            __zfs_make_zvol "${_name}/disk${_num}" "${_size}" "0" "${_zfs_opts}"
            _diskname="disk${_num}"
            ;;
        sparse-zvol)
            __zfs_make_zvol "${_name}/disk${_num}" "${_size}" "1" "${_zfs_opts}"
            _diskname="disk${_num}"
            ;;
        file)
            truncate -s "${_size}" "${vm_dir}/${_name}/disk${_num}.img"
            [ $? -ne 0 ] && __err "failed to create sparse file for disk image"
            _diskname="disk${_num}.img"
            ;;
        *)
            __err "device type must be one of the following: zvol sparse-zvol file"
            ;;
    esac

    # update configuration
    __vm_config_set "${_name}" "disk${_num}_name" "${_diskname}"
    __vm_config_set "${_name}" "disk${_num}_type" "${_emulation}" "1"
    __vm_config_set "${_name}" "disk${_num}_dev" "${_device}" "1"
    [ $? -ne 0 ] && __err "disk image created but errors while updating guest configuration"
}

# add network interface to guest
#
# @private
# @param string _name name of the guest
# @param string _switch the switch name for this interface
#
__vm_add_network(){
    local _name="$1"
    local _switch="$2"
    local _num=0 _curr _emulation

    [ -z "${_switch}" ] && __usage

    __config_load "${vm_dir}/${_name}/${_name}.conf"

    while [ 1 ]; do
        _emulation="${_curr}"
        __config_get "_curr" "network${_num}_type"
        [ -z "${_curr}" ] && break
        _num=$((_num + 1))
    done

    # handle no existing network
    : ${_emulation:=virtio-net}

    # update configuration
    __vm_config_set "${_name}" "network${_num}_type" "${_emulation}"
    __vm_config_set "${_name}" "network${_num}_switch" "${_switch}" "1"
    [ $? -ne 0 ] && __err "errors encountered while updating guest configuration"
}

# 'vm install'
# install os to a virtual machine
#
# @param string _name the guest to install to
# @param string _iso the iso file in $vm_dir/.iso to use
#
__vm_install(){
    local _name="$1"
    local _iso="$2"

    [ -z "${_name}" -o -z "${_iso}" ] && __usage

    # just run start with an iso
    __vm_start "$1" "$2"
}

# 'vm startall'
# start all virtual machines listed in rc.conf:$vm_list
#
__vm_startall(){
    local _vm _done _conf

    [ -z "${vm_list}" ] && exit

    # default to 5 second delay, and don't try starting multiple guests on stdio...
    : ${vm_delay:=5}
    VM_FOREGROUND=""

    for _vm in ${vm_list}; do
        [ -n "${_done}" ] && sleep ${vm_delay}

        # check conf file just to make sure not an obvious typo in rc.conf
        # if the vm is not configured correctly __vm_start will exit
        # causing any further vms to not be started
        _conf="${vm_dir}/${_vm}/${_vm}.conf"

        if [ -f "${_conf}" ]; then
            echo "Starting ${_vm}..."
            __vm_start "${_vm}"
            _done=1
        fi
    done
}

# 'vm stopall'
# stop all bhyve instances
# note this will also stop instances not started by vm-bhyve
#
__vm_stopall(){
    local _pids=$(pgrep -f 'bhyve:')

    echo "Shutting down all bhyve virtual machines"
    killall bhyve
    sleep 1
    killall bhyve
    wait_for_pids ${_pids}
}

# 'vm start'
# start a virtual machine
#
# @param string _name the name of the guest to start
# @param optional string _iso iso file is this is an install (can only be provided through 'vm install' command)
#
__vm_start(){
    local _name="$1"
    local _iso="$2"
    local _conf _host
    local _cpu _memory _disk _guest

    _conf="${vm_dir}/${_name}/${_name}.conf"

    [ -z "${_name}" ] && __usage
    [ ! -f "${_conf}" ] && __err "${_name} does not seem to be a valid virtual machine"

    # confirm we aren't running
    __vm_confirm_stopped "${_name}" "1" || return 1

    # check basic settings before going into background mode
    __config_load "${_conf}"
    __config_get "_guest" "guest"
    __config_get "_cpu" "cpu"
    __config_get "_memory" "memory"
    __config_get "_disk" "disk0_name"

    if [ -z "${_cpu}" -o -z "${_memory}" -o -z "${_disk}" ]; then
        __warn "incomplete virtual machine configuration"
        return 1
    fi

    __vm_tmp_upgrade_check

    # run background process to actually start bhyve
    # this will run as long as vm is running, including restarting bhyve after guest reboot
    if [ -n "${VM_FOREGROUND}" ]; then
        $0 -f _run "${_name}" "${_iso}"
    else
        $0 _run "${_name}" "${_iso}" >/dev/null 2>&1 &
    fi
}

# 'vm stop'
# send a kill signal to the specified guest
#
# @param string _name name of the guest to stop
#
__vm_stop(){
    local _name="$1"
    local _pid _loadpid

    [ -z "${_name}" ] && __usage

    while [ -n "${_name}" ]; do
        if [ ! -e "/dev/vmm/${_name}" ]; then
            __warn "${_name} doesn't appear to be a running virtual machine"
        else
            _pid=$(pgrep -fx "bhyve: ${_name}")
            _loadpid=$(pgrep -fl "grub-bhyve|bhyveload" | grep " ${_name}\$" |cut -d' ' -f1)

            if [ -n "${_pid}" ]; then
                echo "Sending ACPI shutdown to ${_name}"
                kill "${_pid}"
                sleep 1
                kill "${_pid}"
            elif [ -n "${_loadpid}" ]; then
                if __confirm "Guest ${_name} is in bootloader stage, do you wish to force exit"; then
                    echo "Killing ${_name}"
                    kill "${_loadpid}"
                    bhyvectl --destroy --vm=${_name} >/dev/null 2>&1
                fi
            else
                __warn "unable to locate process id for ${_name}"
            fi
        fi
        shift
        _name="$1"
    done
}

# 'vm reset'
# force reset
#
# @param string _name name of the guest
#
__vm_reset(){
    local _name="$1"

    [ -z "${_name}" ] && __usage
    [ ! -e "/dev/vmm/${_name}" ] && __err "${_name} doesn't appear to be a running virtual machine"

    __confirm "Are you sure you want to forcefully reset this virtual machine" || exit 0
    bhyvectl --force-reset --vm=${_name}
}

# 'vm poweroff'
# force poweroff
#
# @param string _name name of the guest
#
__vm_poweroff(){
    local _name="$1"

    [ -z "${_name}" ] && __usage
    [ ! -e "/dev/vmm/${_name}" ] && __err "${_name} doesn't appear to be a running virtual machine"

    __confirm "Are you sure you want to forcefully poweroff this virtual machine" || exit 0
    bhyvectl --force-poweroff --vm=${_name}
}

# 'vm destroy'
# completely remove a guest
#
# @param string _name name of the guest
#
__vm_destroy(){
    local _name="$1"

    [ -z "${_name}" ] && __usage
    [ ! -e "${vm_dir}/${_name}" ] && __err "${_name} doesn't appear to be a valid virtual machine"

    # make sure it's stopped!
    __vm_confirm_stopped "${_name}" || exit 1

    __confirm "Are you sure you want to completely remove this virtual machine" || exit 0
    __zfs_destroy_dataset "${_name}"
    [ -e "${vm_dir}/${_name}" ] && rm -R "${vm_dir}/${_name}"
}

# 'vm rename'
# rename an existing guest
#
# @param string _old the existing guest name
# @param string _new the new guest name
#
__vm_rename(){
    local _old="$1"
    local _new="$2"

    [ -z "${_old}" -o -z "${_new}" ] && __usage
    [ ! -e "${vm_dir}/${_old}/${_old}.conf" ] && __err "${_old} doesn't appear to be a valid virtual machine"
    [ -d "${vm_dir}/${_new}" ] && __err "directory ${vm_dir}/${_new} already exists"

    __vm_check_name "${_new}" || __err "invalid virtual machine name - '${_name}'"

    # confirm guest stopped
    __vm_confirm_stopped "${_old}" || exit 1

    # rename zfs dataset
    __zfs_rename_dataset "${_old}" "${_new}"

    # rename folder if it still exists (shouldn't if zfs mode and rename worked)
    if [ -d "${vm_dir}/${_old}" ]; then
        mv "${vm_dir}/${_old}" "${vm_dir}/${_new}" >/dev/null 2>&1
        [ $? -ne 0 ] && __err "failed to rename guest directory"
    fi

    # rename config file
    mv "${vm_dir}/${_new}/${_old}.conf" "${vm_dir}/${_new}/${_new}.conf" >/dev/null 2>&1
    [ $? -ne 0 ] && __err "changed guest directory but failed to rename configuration file"
}

# 'vm console'
# connect to the console (nmdm) of the specified guest
# we store the nmdm for com1 & com2 in $vm_dir/{guest}/console
# if no port is specified, we use the first one that is specified in the configuration file
# so if comports="com2 com1", it will connect to com2
# the boot loader always using the nmdm device of the first com port listed
#
# @param string _name name of the guest
# @param string _port the port to connect to (default = first in configuration)
#
__vm_console(){
    local _name="$1"
    local _port="$2"
    local _console

    [ -z "${_name}" ] && __usage
    [ ! -e "/dev/vmm/${_name}" ] && __err "${_name} doesn't appear to be a running virtual machine"
    [ ! -e "${vm_dir}/${_name}/console" ] && __err "can't locate console data for ${_name}"

    # did user specify a com port?
    # if not, get first in the file (the first will also be the console used for loader)
    if [ -n "${_port}" ]; then
        _console=$(grep "${_port}=" "${vm_dir}/${_name}/console" | cut -d= -f2)
    else
        _console=$(head -n 1 "${vm_dir}/${_name}/console" | cut -d= -f2)
    fi

    [ -z "${_console}" ] && __err "unable to locate console device for this virtual machine"
    cu -l "${_console}"
}

# 'vm configure'
# configure a machine (edit the configuration file)
#
# @param string _name name of the guest
#
__vm_configure(){
    local _name="$1"

    [ -z "${_name}" ] && __usage
    [ -z "${EDITOR}" ] && EDITOR=vi
    [ ! -e "${vm_dir}/${_name}/${_name}.conf" ] && \
        __err "cannot locate configuration file for virtual machine: ${_name}"

    $EDITOR "${vm_dir}/${_name}/${_name}.conf"
}

# 'vm iso'
# list iso images or get a new one
#
# @param string _url if specified, the url will be fetch'ed into $vm_dir/.iso
#
__vm_iso(){
    local _url="$1"

    if [ -n "${_url}" ]; then
        fetch -o "${vm_dir}/.iso" "${_url}"
    else
        echo "FILENAME"
        ls -1 "${vm_dir}/.iso"
    fi
}

# 'vm passthru'
# show a list of available passthrough devices
# and their device number
#
__vm_passthru(){
    local _dev _sbf _desc _ready
    local _format="%-10s %-12s %-12s %s\n"

    printf "${_format}" "DEVICE" "BHYVE ID" "READY" "DESCRIPTION"

    pciconf -l | awk -F'[@:]' '{ print $1,$3 "/" $4 "/" $5}' | \
    while read _dev _sbf; do

        _ready=$(echo "${_dev}" | grep ^ppt)
        [ -n "${_ready}" ] && _ready="Yes"

        _desc=$(pciconf -lv | grep -A2 "^${_dev}@" | tail -n1 | grep device | cut -d\' -f2)
        printf "${_format}" "${_dev}" "${_sbf}" "${_ready:-No}" "${_desc:--}"
    done
}

# set a value in vm configuration file
# we check for newline at the end as sysrc won't add it
# and that will mess up the new key and the existing one
# from the end of the file
#
# @param string _name guest name
# @param string _key config key to set
# @param string _value value
# @param int _skip_newline_check skip the check for newline
#
__vm_config_set(){
    local _name="$1"
    local _key="$2"
    local _value="$3"
    local _skip_newline_check="$4"
    local _newline

    if [ -z "${_skip_newline_check}" ]; then
        _newline=$(tail -1 "${vm_dir}/${_name}/${_name}.conf" | wc -l | tr -d " ")
        [ "${_newline}" -eq "0" ] && echo "" >> "${vm_dir}/${_name}/${_name}.conf"
    fi

    sysrc -inqf "${vm_dir}/${_name}/${_name}.conf" "${_key}=${_value}" >/dev/null 2>&1
    return $?
}

# temporary check for guests that need reconfiguring
#
__vm_tmp_upgrade_check(){

    # hopefully no guest means it's a recent guest
    [ -z "${_guest}" ] && return 0

    case "${_guest}" in
        freebsd) ;&
        linux) ;&
        generic) ;&
        windows) ;;
        *)
            echo "******"
            echo "  This guest needs to be updated to use newer grub-bhyve support"
            echo "  Please see the GitHub readme for more information"
            echo "******"
            exit 1
            ;;
    esac
}
