#!/bin/sh

__update () {
  local _name _dataset _fulluuid _mountpoint _date _jail_type _non_interactive \
        _resolv _jail_release

  _name=$2
  _resolv="0"

    # Start looking for any command switches, if nothing special is found
    # then proceed with the normal update procedure
    case "$1" in
        -n)
            _dataset=$(__find_jail ${_name}) || exit $?
            _fulluuid="$(__check_name ${_name} ${_dataset})"
            _mountpoint="$(__get_jail_prop mountpoint ${_fulluuid} ${_dataset})"
            _jail_type="$(__get_jail_prop type ${_fulluuid} ${_dataset})"

            # If a git jail is found, update the remote and show the user the
            # status output
            case ${_jail_type} in
                gitjail)
                    git -C ${_mountpoint}/_ remote update
                    git -C ${_mountpoint}/_ status -uno
                    ;;
            esac
            ;;
        -P|-p)
            if [ ! -e "${iocroot}/releases/${_name}/root" ] ; then
                __die "${_name} does not exist!"
            fi

            _jail_release="$(__get_jail_prop release ${_name} ${_dataset})"

            # The user might be supplying a RELEASE instead of a jail
            if [ "${_jail_release}" = "none" ] ; then
                _jail_release="${_name}"
            fi

            mount -t devfs devfs ${iocroot}/releases/${_name}/root/dev
            cat /etc/resolv.conf > \
                ${iocroot}/releases/${_name}/root/etc/resolv.conf

            # Lets us make sure any templates using 9.3-RELEASE get the right
            # treatment for non-interactive portsnap.
            case "${_jail_release}" in
                9.3-RELEASE)
                    # Remove the interactive check in portsnap
                    sed 's/\[ ! -t 0 \]/false/' \
                        ${iocroot}/releases/${_name}/root/usr/sbin/portsnap \
                        > ${iocroot}/releases/${_name}/root/tmp/portsnap
                    chmod +x ${iocroot}/releases/${_name}/root/tmp/portsnap
                    chroot ${iocroot}/releases/${_name}/root \
                        env UNAME_r="${_name}" env PAGER="/bin/cat" \
                        /tmp/portsnap fetch update
                    rm ${iocroot}/releases/${_name}/root/tmp/portsnap
                    ;;
                *)
                    chroot ${iocroot}/releases/${_name}/root \
                        env UNAME_r="${_name}" env PAGER="/bin/cat" \
                        portsnap --interactive fetch update
                    ;;
                esac

                umount ${iocroot}/releases/${_name}/root/dev
                rm ${iocroot}/releases/${_name}/root/etc/resolv.conf
                ;;
        *)
            shift 3
            _name=$1
            _dataset=$(__find_jail ${_name}) || exit $?
            _fulluuid="$(__check_name ${_name} ${_dataset})"
            _mountpoint="$(__get_jail_prop mountpoint ${_fulluuid} ${_dataset})"
            _date=$(date "+%F_%T")
            _jail_type="$(__get_jail_prop type ${_fulluuid} ${_dataset})"
            _jail_release="$(__get_jail_prop release ${_fulluuid} ${_dataset})"

            # Check all jail types and do the right thing
            case ${_jail_type} in
                basejail)
                    if [ -e "${iocroot}/releases/${_jail_release}/root/etc/freebsd-update.conf" ] ; then
                        if [ ! -e ${iocroot}/releases/${_jail_release}/root/etc/resolv.conf ] ; then
                            export _resolv="1"
                            cat /etc/resolv.conf > \
                            ${iocroot}/releases/${_jail_release}/root/etc/resolv.conf
                        fi

                        mount -t devfs devfs \
                            ${iocroot}/releases/${_jail_release}/root/dev

                        case "${_jail_release}" in
                            9.3-RELEASE | 10.1-RELEASE)
                                # Remove the interactive check in freebsd-update
                                sed 's/\[ ! -t 0 \]/false/' \
                                    ${iocroot}/releases/${_jail_release}/root/usr/sbin/freebsd-update \
                                    > ${iocroot}/releases/${_jail_release}/root/tmp/freebsd-update
                                chmod +x ${iocroot}/releases/${_jail_release}/root/tmp/freebsd-update
                                chroot ${iocroot}/releases/${_jail_release}/root \
                                    env UNAME_r="${_jail_release}" \
                                    env PAGER="/bin/cat" \
                                    /tmp/freebsd-update fetch
                                rm ${iocroot}/releases/${_jail_release}/root/tmp/freebsd-update
                                ;;
                            *)
                                chroot ${iocroot}/releases/${_jail_release}/root \
                                    env UNAME_r="${_jail_release}" \
                                    env PAGER="/bin/cat" \
                                    freebsd-update --not-running-from-cron fetch
                                ;;
                            esac

                        chroot ${iocroot}/releases/${_jail_release}/root \
                            env UNAME_r="${_jail_release}" env PAGER="/bin/cat" \
                            freebsd-update install
                        umount ${iocroot}/releases/${_jail_release}/root/dev

                        if [ "${_resolv}" = "1" ] ; then
                            if [ $? -eq 0 ] ; then
                                rm -f \
                                    ${iocroot}/releases/${_jail_release}/root/etc/resolv.conf
                            fi
                        fi
                    fi

                    # Re-clone required filesystems
                    __check_basejail "${_fulluuid}" "${_dataset}"
                    __umount_basejail "${_fulluuid}" "${_dataset}"
                    echo "* ${_name} updated successfully"
                    ;;
                gitjail)
                    # Update the local git repo to the latest remote
                    git -C ${_mountpoint}/root pull
                    ;;
                *)
                    echo "* Creating back-out update snapshot."
                    __snapshot ${_fulluuid}@ioc-update_${_date}

                    # Override the release that freebsd-update uses and
                    # set the pager to cat so that it isn't interactive
                    echo "* Updating jail.."

                    if [ -e "${iocroot}/releases/${_jail_release}/root/etc/freebsd-update.conf" ] ; then
                        if [ ! -e "${_mountpoint}/root/etc/resolv.conf" ] ; then
                            export _resolv="1"
                            cat /etc/resolv.conf > \
                            ${_mountpoint}/root/etc/resolv.conf
                        fi

                        mount -t devfs devfs \
                            ${_mountpoint}/root/dev

                        case "${_jail_release}" in
                            9.3-RELEASE | 10.1-RELEASE)
                                # Remove the interactive check in freebsd-update
                                sed 's/\[ ! -t 0 \]/false/' \
                                    ${_mountpoint}/root/usr/sbin/freebsd-update \
                                    > ${_mountpoint}/root/tmp/freebsd-update
                                chmod +x ${_mountpoint}/root/tmp/freebsd-update
                                chroot ${_mountpoint}/root \
                                    env UNAME_r="${_jail_release}" \
                                    env PAGER="/bin/cat" \
                                    /tmp/freebsd-update fetch
                                rm ${_mountpoint}/root/tmp/freebsd-update
                                ;;
                            *)
                                chroot ${_mountpoint}/root \
                                    env UNAME_r="${_jail_release}" \
                                    env PAGER="/bin/cat" \
                                    freebsd-update --not-running-from-cron fetch
                                ;;
                            esac

                        chroot ${_mountpoint}/root \
                            env UNAME_r="${_jail_release}" env PAGER="/bin/cat" \
                            freebsd-update install
                        umount ${_mountpoint}/root/dev

                        if [ "${_resolv}" == "1" ] ; then
                            if [ $? -eq 0 ] ; then
                                rm -f ${_mountpoint}/root/etc/resolv.conf
                            fi
                        fi
                    fi

                    echo "* ${_name} updated successfully. Please reboot jail and inspect."
                    echo "* Remove:"
                    echo "    snapshot: ${_dataset}/root@ioc-update_${_date}"
                    echo "* if everything is OK."
                    ;;
            esac

            if [ -z "${_name}" ] ; then
                __die "missing UUID!"
            fi

            if [ -z "${_dataset}" ] ; then
                __die "${_name} not found!"
            fi
            ;;
      esac
}

__upgrade () {
    local _name _release _fulluuid _jails _jail_type _mountpoint _date \
          _oldrelease _resolv _origin _date _etcsuccess

    _name="$1"
    _release="$(uname -r|cut -f 1,2 -d'-'|sed -E 's/CURRENT.*|STABLE.*/RELEASE/')"
    _dataset="$(__find_jail ${_name})" || exit $?
    _fulluuid="$(__check_name ${_name} ${_dataset})"
    _jail_type="$(__get_jail_prop type ${_fulluuid} ${_dataset})"
    _mountpoint="$(__get_jail_prop mountpoint ${_fulluuid} ${_dataset})"
    _oldrelease="$(__get_jail_prop release ${_fulluuid} ${_dataset})"
    _date=$(date "+%F")
    _resolv="0"

    if [ -z "${_name}" ] ; then
        __die "missing UUID!"
    fi

    if [ -z "${_dataset}" ] ; then
        __die "${_name} not found!"
    fi

    if [ -n "${iocset_release}" ] ; then
        _release="${iocset_release}"
    fi

    if [ ! -d "${iocroot}/releases/${_release}" ] ; then
        __error "${_release} not found!"
        echo "  Please run iocell fetch first." >&2
        exit 1
    fi

    # If jail is a basejail then set the new release and check the datasets
    # back up /etc just in case the merge fails
    if [ "${_jail_type}" = "basejail" ] ; then
        _jails=$(__find_jail ALL | \
            awk 'BEGIN { FS = "/" } ; { if (length($NF)==36) { print }}')
        # Check to see if the jail has a stale upgrade snapshot.
        for _fs in ${_jails} ; do
            _snapcheck="$(zfs list -Ht snap | \
                     egrep "${_fulluuid}.*ioc-upgrade" | \
                     awk '{print $1}')"
            if [ ! -z ${_snapcheck} ] ; then
                __error "${_name} has an existing upgrade snapshot" >&2
                printf "%1s ${_snapcheck} \n" >&2
                printf "%1s please remove this before trying to upgrade. \n" >&2
                exit 1
            fi
        done

        __set_jail_prop release="${_release}" "${_fulluuid}" "${_dataset}"
        __check_basejail ${_fulluuid} > /dev/null 2>&1
        cp -Rp ${_mountpoint}/root/etc \
               ${_mountpoint}/etc.old
        zfs snapshot ${_dataset}/root@ioc-upgrade_${_date}
        chroot ${_mountpoint}/root etcupdate -F
        chroot ${_mountpoint}/root etcupdate resolve
        _etcsuccess="$?"

        if [ -d "${_mountpoint}/root/var/spool/mqueue" ] ; then
            chroot ${_mountpoint}/root newaliases > /dev/null 2>&1
        else
            mkdir -p ${_mountpoint}/root/var/spool/mqueue
            chroot ${_mountpoint}/root newaliases > /dev/null 2>&1
        fi

        # Check for zero exit code(success) otherwise bail and revert everything
        if [ "${_etcsuccess}" -eq 0 ] ; then
            echo "* ${_name} upgraded successfully. Please reboot jail and inspect."
            echo "* Remove:"
            echo "    /etc backup: ${_mountpoint}/etc.old"
            echo "    snapshot: ${_dataset}/root@ioc-upgrade_${_date}"
            echo "* if everything is OK."
            __umount_basejail "${_fulluuid}" "${_dataset}"
            exit 0
        else
            echo "" >&2
            echo "* Upgrade failed! Backing out." >&2
            __set_jail_prop release="${_oldrelease}" "${_fulluuid}" \
                "${_dataset}"
            zfs rollback -rRf ${_dataset}/root@ioc-upgrade_${_date}
            zfs destroy ${_dataset}/root@ioc-upgrade_${_date}
            rm -rf ${_mountpoint}/root/etc
            mv ${_mountpoint}/etc.old \
               ${_mountpoint}/root/etc
            __check_basejail "${_fulluuid}"
            __umount_basejail "${_fulluuid}" "${_dataset}"
            exit 1
        fi
    elif [ "${_jail_type}" = "thickjail" ] ; then
        # __spinner needs to be recreated here as we redirect with __update
        _spinner='/-\|'

        printf "* Updating ${_name} to latest patch level first:  "

        while true; do
            printf '\b%.1s' "${_spinner}"
            _spinner=${_spinner#?}${_spinner%???}
            sleep .1
        done &

        trap "kill $!" 2 3
        __update ${_fulluuid} > /dev/null 2>&1
        printf "\b%1s\n" "done!" ; kill $! && trap " " 2 3

        echo "* Creating back-out upgrade snapshot."
        zfs snapshot ${_dataset}/root@ioc-upgrade_${_date}
        echo "* Upgrading jail."
        if [ -e "${iocroot}/releases/${_release}/root/etc/freebsd-update.conf" ] ; then
            if [ ! -e "${_mountpoint}/root/etc/resolv.conf" ] ; then
                export _resolv="1"
                cat /etc/resolv.conf > \
                    ${_mountpoint}/root/etc/resolv.conf
            fi
            mount -t devfs devfs \
                ${_mountpoint}/root/dev
            chroot ${_mountpoint}/root \
                env UNAME_r="${_oldrelease}" env PAGER="/bin/cat" \
                freebsd-update -r ${_release} upgrade
            if [ $? -eq 0 ] ; then
                while [ $? -eq 0 ] ; do
                    chroot ${_mountpoint}/root \
                        env UNAME_r="${_oldrelease}" env PAGER="/bin/cat" \
                        freebsd-update -r ${_release} install
                done
                if [ "${_resolv}" == "1" ] ; then
                    rm -f ${_mountpoint}/root/etc/resolv.conf
                fi
                umount ${_mountpoint}/root/dev
                # Set jail's zfs property to new release
                __set_jail_prop release="${_release}" "${_fulluuid}" \
                    "${_dataset}"
            else
                echo "* Upgrade failed, aborting install." >&2
                zfs rollback -rRf ${_dataset}/root@ioc-upgrade_${_date}
                zfs destroy ${_dataset}/root@ioc-upgrade_${_date}
                exit 1
            fi
        fi

        echo "* ${_name} upgraded successfully. Please reboot jail and inspect."
        echo "* Remove:"
        echo "    snapshot: ${_dataset}/root@ioc-upgrade_${_date}"
        echo "* if everything is OK."
    else
        __die "upgrade is not supported for this type of jail"
    fi
}
