#!/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
#
core::list(){
    local _name _loader _cpu _our_host
    local _memory _run _vm _auto _num _uefi _graphics _vnc
    local _format="%-15s %-15s %-11s %-6s %-9s %-20s %-12s %s\n"

    _our_host=$(hostname)

    vm::running_load
    printf "${_format}" "NAME" "DATASTORE" "LOADER" "CPU" "MEMORY" "VNC" "AUTOSTART" "STATE"

    for _ds in ${VM_DATASTORE_LIST}; do
        datastore::get "${_ds}" || continue

        ls -1 "${VM_DS_PATH}" 2>/dev/null | \
        while read _name; do
            [ ! -e "${VM_DS_PATH}/${_name}/${_name}.conf" ] && continue

            config::load "${VM_DS_PATH}/${_name}/${_name}.conf"
            config::get "_loader" "loader" "none"
            config::get "_cpu" "cpu"
            config::get "_memory" "memory"
            config::get "_uefi" "uefi"
            config::get "_graphics" "graphics"

            # show uefi in loader column if it's a uefi guest
            util::checkyesno "${_uefi}" && _loader="uefi"

            # defaults
            _vnc="-"

            # check if the guest is running
            if vm::running_check "_run" "${_name}" || \
               [ -e "${VM_DS_PATH}/${_name}/run.lock" -a "$(head -n1 ${VM_DS_PATH}/${_name}/run.lock 2>/dev/null)" = "${_our_host}" ]; then

                # if running and graphics, try to get vnc port
                if util::checkyesno "${_graphics}"; then
                    _vnc=$(grep vnc "${VM_DS_PATH}/${_name}/console" 2>/dev/null |cut -d= -f2)
                    [ -z "${_vnc}" ] && _vnc="-"
                fi
            fi

            _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_DS_PATH}/${_name}/run.lock" ]; then
                _run=$(head -n1 "${VM_DS_PATH}/${_name}/run.lock")
                _run="Locked (${_run})"
            fi

            printf "${_format}" "${_name}" "${_ds}" "${_loader}" "${_cpu}" "${_memory}" "${_vnc}" "${_auto}" "${_run}"
        done
    done
}

# 'vm check name'
# check name of virtual machine
#
# @return int 0 if name is valid
#
core::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
#
core::create(){
    local _name _opt _size _vmdir _disk _disk_dev _uuid _num=0
    local _zfs_opts _disk_size _template="default" _ds="default" _ds_path

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

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

    [ -z "${_name}" ] && util::usage

    # check guest name
    core::check_name "${_name}" || util::err "invalid virtual machine name - '${_name}'"
    datastore::get_guest "${_name}" && util::err "virtual machine already exists in ${VM_DS_PATH}/${_name}"
    datastore::get "${_ds}" || util::err "unable to load datastore - '${_ds}'"

    [ ! -f "${vm_dir}/.templates/${_template}.conf" ] && \
        util::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}" ] && util::err "template is missing disk0_name specification"

    # if we're on zfs, make a new filesystem
    zfs::make_dataset "${VM_DS_ZFS_DATASET}/${_name}" "${_zfs_opts}"

    [ ! -d "${VM_DS_PATH}/${_name}" ] && mkdir "${VM_DS_PATH}/${_name}" >/dev/null 2>&1
    [ ! -d "${VM_DS_PATH}/${_name}" ] && util::err "unable to create virtual machine directory ${VM_DS_PATH}/${_name}"

    cp "${vm_dir}/.templates/${_template}.conf" "${VM_DS_PATH}/${_name}/${_name}.conf"
    [ $? -ne 0 ] && util::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)
    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 "${VM_DS_ZFS_DATASET}/${_name}/${_disk}" "${_disk_size}" "0" "${_zfs_opts}"
                ;;			
            sparse-zvol)
                zfs::make_zvol "${VM_DS_ZFS_DATASET}/${_name}/${_disk}" "${_disk_size}" "1" "${_zfs_opts}"
                ;;
            *)
                truncate -s "${_disk_size}" "${VM_DS_PATH}/${_name}/${_disk}"
                [ $? -ne 0 ] && util::err "failed to create sparse file for disk image"
                ;;
        esac

        # scrap size option from guest template
        sysrc -inxqf "${VM_DS_PATH}/${_name}/${_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
#
core::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} ;;
            *) util::usage ;;
        esac
    done

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

    # check guest
    [ -z "${_name}" ] && util::usage
    datastore::get_guest "${_name}" || "${_name} does not appear to be a valid virtual machine"

    case "${_device}" in
        disk)    core::add_disk "${_name}" "${_type}" "${_sopt}" ;;
        network) core::add_network "${_name}" "${_sopt}" ;;
        *)       util::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
#
core::add_disk(){
    local _name="$1"
    local _device="$2"
    local _size="$3"
    local _num=0 _curr _diskname _emulation _zfs_opts

    : ${_device:=file}

    [ -z "${_size}" ] && util::usage

    # get the last existing disk
    config::load "${VM_DS_PATH}/${_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}" ] && util::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 "${VM_DS_ZFS_DATASET}/${_name}/disk${_num}" "${_size}" "0" "${_zfs_opts}"
            _diskname="disk${_num}"
            ;;
        sparse-zvol)
            zfs::make_zvol "${VM_DS_ZFS_DATASET}/${_name}/disk${_num}" "${_size}" "1" "${_zfs_opts}"
            _diskname="disk${_num}"
            ;;
        file)
            truncate -s "${_size}" "${VM_DS_PATH}/${_name}/disk${_num}.img"
            [ $? -ne 0 ] && util::err "failed to create sparse file for disk image"
            _diskname="disk${_num}.img"
            ;;
        *)
            util::err "device type must be one of the following: zvol sparse-zvol file"
            ;;
    esac

    # update configuration
    config::set "${_name}" "disk${_num}_name" "${_diskname}"
    config::set "${_name}" "disk${_num}_type" "${_emulation}" "1"
    config::set "${_name}" "disk${_num}_dev" "${_device}" "1"
    [ $? -ne 0 ] && util::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
#
core::add_network(){
    local _name="$1"
    local _switch="$2"
    local _num=0 _curr _emulation

    [ -z "${_switch}" ] && util::usage

    config::load "${VM_DS_PATH}/${_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
    config::set "${_name}" "network${_num}_type" "${_emulation}"
    config::set "${_name}" "network${_num}_switch" "${_switch}" "1"
    [ $? -ne 0 ] && util::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
#
core::install(){
    local _name="$1"
    local _iso="$2"

    [ -z "${_name}" -o -z "${_iso}" ] && util::usage
    [ ! -e "${vm_dir}/.iso/${_iso}" ] && util::err "${_iso} does not appear to be a valid iso file"

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

# 'vm startall'
# start all virtual machines listed in rc.conf:$vm_list
#
core::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}

        echo "Starting ${_vm}..."
        core::start "${_vm}"
        _done=1
    done
}

# 'vm stopall'
# stop all bhyve instances
# note this will also stop instances not started by vm-bhyve
#
core::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)
#
core::start(){
    local _name="$1"
    local _iso="$2"
    local _cpu _memory _disk _guest _loader _console _tmux_cmd _tmux_name

    [ -z "${_name}" ] && util::usage

    # try to find guest
    if ! datastore::get_guest "${_name}"; then
        util::warn "${_name} does not seem to be a valid virtual machine"
        return 1
    fi

    # confirm we aren't running
    vm::confirm_stopped "${_name}" "1" || return 1

    # check basic settings before going into background mode
    config::load "${VM_DS_PATH}/${_name}/${_name}.conf"
    config::get "_cpu" "cpu"
    config::get "_memory" "memory"
    config::get "_disk" "disk0_name"
    config::get "_loader" "loader"

    # check minimum configuration
    if [ -z "${_cpu}" -o -z "${_memory}" -o -z "${_disk}" ]; then
        util::warn "incomplete virtual machine configuration"
        return 1
    fi

    # we can only load freebsd without unrestricted guest support
    if [ -n "${VM_NO_UG}" -a "${_loader}" != "bhyveload" ]; then
        util::warn "no unrestricted guest support in cpu. only single vcpu FreeBSD guests supported"
        return 1
    fi

    # check for tmux
    config::core_get "_console" "console" "nmdm"
    _tmux_cmd=$(which tmux)

    if [ "${_console}" = "tmux" -a -z "${_tmux_cmd}" ]; then
        util::warn "tmux support enabled but misc/tmux not found"
        return 1
    fi

    # 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}"
    elif [ "${_console}" = "tmux" ]; then
        # can't have dots in tmux session :( (looks like it may use . to separate window.pane)
        # use ~ which we don't normally allow
        _tmux_name=$(echo "${_name}" | tr "." "~")

        ${_tmux_cmd} new -ds "${_tmux_name}" $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
#
core::stop(){
    local _name="$1"
    local _pid _loadpid

    [ -z "${_name}" ] && util::usage

    while [ -n "${_name}" ]; do
        if [ ! -e "/dev/vmm/${_name}" ]; then
            util::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 util::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
                util::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
#
core::reset(){
    local _name="$1"

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

    util::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
#
core::poweroff(){
    local _name="$1"

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

    util::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
#
core::destroy(){
    local _name="$1"

    [ -z "${_name}" ] && util::usage
    datastore::get_guest "${_name}" || util::err "${_name} doesn't appear to be a valid virtual machine"

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

    util::confirm "Are you sure you want to completely remove this virtual machine" || exit 0
    zfs::destroy_dataset "${VM_DS_ZFS_DATASET}/${_name}"
    [ -e "${VM_DS_PATH}/${_name}" ] && rm -R "${VM_DS_PATH}/${_name}"
}

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

    [ -z "${_old}" -o -z "${_new}" ] && util::usage
    core::check_name "${_new}" || util::err "invalid virtual machine name - '${_name}'"

    datastore::get_guest "${_new}" && util::err "directory ${VM_DS_PATH}/${_new} already exists"
    datastore::get_guest "${_old}" || util::err "${_old} doesn't appear to be a valid virtual machine"

    # 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_DS_PATH}/${_old}" ]; then
        mv "${VM_DS_PATH}/${_old}" "${VM_DS_PATH}/${_new}" >/dev/null 2>&1
        [ $? -ne 0 ] && util::err "failed to rename guest directory"
    fi

    # rename config file
    mv "${VM_DS_PATH}/${_new}/${_old}.conf" "${VM_DS_PATH}/${_new}/${_new}.conf" >/dev/null 2>&1
    [ $? -ne 0 ] && util::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)
#
core::console(){
    local _name="$1"
    local _port="$2"
    local _console _tmux _tmux_cmd _tmux_name

    [ -z "${_name}" ] && util::usage

    datastore::get_guest "${_name}" || util::err "${_name} doesn't appear to be a valid virtual machine"
    [ ! -e "/dev/vmm/${_name}" ] && util::err "${_name} doesn't appear to be a running virtual machine"

    if [ -e "${VM_DS_PATH}/${_name}/console" ]; then

        # 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_DS_PATH}/${_name}/console" | cut -d= -f2)
        else
            _console=$(head -n 1 "${VM_DS_PATH}/${_name}/console" | grep "^com" | cut -d= -f2)
        fi
    fi

    # check tmux if we don't have a console
    if [ -z "${_console}" ]; then
        _tmux_cmd=$(which tmux)

        if [ -n "${_tmux_cmd}" ]; then
            _tmux=$("${_tmux_cmd}" ls |grep "^${_name}:")

            if [ -n "${_tmux}" ]; then
                _tmux_name=$(echo "${_name}" | tr "." "~")

                ${_tmux_cmd} attach -t ${_tmux_name}
                exit
            fi
        fi
    fi

    [ -z "${_console}" ] && util::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
#
core::configure(){
    local _name="$1"

    [ -z "${_name}" ] && util::usage
    [ -z "${EDITOR}" ] && EDITOR=vi

    datastore::get_guest "${_name}" || \
        util::err "cannot locate configuration file for virtual machine: ${_name}"

    $EDITOR "${VM_DS_PATH}/${_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
#
core::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
#
core::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
}
