#!/bin/sh
#
# Copyright (c) 2017-2018 Aaron LI <aly@aaronly.me>
# MIT License
#
# Tool to update a DragonFly BSD system using binary releases or
# snapshot builds.
#

# Exit immediately if any untested command fails in non-interactive mode
set -e
# Exit with an error when an undefined variable is referenced
set -u

DEBUG=${DEBUG:-""}

VERSION="0.1.3"
NAME="dfly-update"
TOOLDIR="${0%/*}"
PREFIX="${PREFIX:-${TOOLDIR}}"
CONFIGFILE="${PREFIX}/${NAME}.conf"

#
# Error Codes
#

EC_OS=10
EC_TMPFILE=11
EC_MD5=12
EC_ARGS=13
EC_FETCH=14
EC_MOUNT=15
EC_UMOUNT=16
EC_VN=17
EC_TAR=18
EC_MTREE=19
EC_CPDUP=20
EC_NOFILE=21


#
# Default Configurations
#

# Base URL to Remote DragonFly BSD images
URL_BASE="https://mirror-master.dragonflybsd.org"
URL_MASTER="${URL_BASE}/snapshots/x86_64/images"
URL_RELEASE="${URL_BASE}/iso-images"

# Default to track the same branch as the installed system
# * MASTER (i.e., the DEVELOPMENT branch)
# * RELEASE
# * (empty) - same as the local installed branch
UPDATE_BRANCH=

# Temporary directory to cache the image, etc, ...
CACHE_DIR="/var/tmp/${NAME}"
# Directory to mount the system image
MNT_DIR="/mnt/${NAME}"
# Backup directory
BACK_DIR="/var/backups/${NAME}"

# Ignored files to be kept from overriding by the upgrade.
FILES_IGNORE='
    /boot/loader.conf
    /etc/crypttab
    /etc/fstab
    /etc/group
    /etc/localtime
    /etc/master.passwd
    /etc/passwd
    /etc/pwd.db
    /etc/rc.conf
    /etc/spwd.db
'
# Filename suffix for the temporarily backup files to avoid overriding.
BAK_SUF="__bak__"
# Filename suffix for newly installed config files that need (manual)
# merge.
NEW_SUF="__new__"


#
# XXX: Variables shared across two or more functions
#
_FILENAME=  # The latest remote system image filename
_MD5=
_HAS_UPDATE=


#
# Helper Functions
#

debug() {
    # Add "|| true" to work with "set -e"
    [ -n "${DEBUG}" ] && echo "DEBUG: $@" >&2 || true
}

log() {
    echo "$@" >&2
}

warn() {
    echo "WARNING: $@" >&2
}

error() {
    local ec="$1"
    shift
    echo "ERROR: $@" >&2
    exit ${ec}
}

check_os() {
    [ $# -eq 0 ] || \
        error ${EC_ARGS} "check_os: invalid arguments: $@"
    local os=$(uname -s)
    if [ "${os}" != "DragonFly" ]; then
        error ${EC_OS} "Not a DragonFly BSD system!"
    fi
}

# contains(string, substring)
#
# Returns 0 if the specified string contains the specified substring,
# otherwise returns 1.
#
# Credit: https://stackoverflow.com/a/8811800
contains() {
    [ $# -eq 2 ] || \
        error ${EC_ARGS} "contains: invalid arguments: $@"
    local string="$1"
    local substring="$2"
    if [ "${string#*$substring}" != "$string" ]; then
        return 0  # $substring is in $string
    else
        return 1  # $substring is not in $string
    fi
}

# Determine the branch from the image filename
get_branch_filename() {
    [ $# -eq 1 ] || \
        error ${EC_ARGS} "get_branch_filename: invalid arguments: $@"
    if contains "$1" "-DEV-"; then
        echo "MASTER"
    else
        echo "RELEASE"
    fi
}

# Determine whether the given name refers to the master branch?
is_master_branch() {
    [ $# -eq 1 ] || \
        error ${EC_ARGS} "is_master_branch: invalid arguments: $@"
    case "$1" in
        master|MASTER|[dD][eE][vV]*)
            return 0
            ;;
        *)
            return 1
            ;;
    esac
}

# Get the branch of the installed system
# * DEVELOPMENT
# * RELEASE
get_local_branch() {
    [ $# -eq 0 ] || \
        error ${EC_ARGS} "get_local_branch: invalid arguments: $@"
    check_os
    uname -r | awk -F'-' '{ print $2 }'
}

# Get the version of local installed system
get_local_version() {
    [ $# -eq 0 ] || \
        error ${EC_ARGS} "get_local_version: invalid arguments: $@"
    check_os
    local version=$(uname -v | awk '{ print $2 }')
    echo "${version}" | awk -F'-' '{ print $1 }' | tr -d 'v'
}

# Get the URL of the MD5 list
get_md5list_url() {
    [ $# -eq 1 ] || \
        error ${EC_ARGS} "get_md5list_url: invalid arguments: $@"
    local branch="$1"
    if is_master_branch "${branch}"; then
        echo "${URL_MASTER}/CHECKSUM.MD5"
    else
        echo "${URL_RELEASE}/md5.txt"
    fi
}

# Determine the URL of the given image filename
get_image_url() {
    [ $# -eq 1 ] || \
        error ${EC_ARGS} "get_image_url: invalid arguments: $@"
    local filename="$1"
    local branch=$(get_branch_filename ${filename})
    if is_master_branch "${branch}"; then
        echo "${URL_MASTER}/${filename}"
    else
        echo "${URL_RELEASE}/${filename}"
    fi
}

# Get the latest remote system image
# Returns:
# "_filename='<latest.iso/img>'; _md5='<md5/of/latest.iso/img>'"
get_latest_image() {
    [ $# -eq 1 ] || \
        error ${EC_ARGS} "get_latest_image: invalid arguments: $@"
    local branch="$1"
    local url_checksum=$(get_md5list_url ${branch})
    local tmpchecksum=$(mktemp -t ${NAME}) || exit ${EC_TMPFILE}
    local latest_filename latest_md5 line
    log "Fetch remote systems checksum: ${url_checksum}"
    fetch -q -o ${tmpchecksum} "${url_checksum}" || exit ${EC_FETCH}
    if is_master_branch "${branch}"; then
        line=$(fgrep '.img.bz2' ${tmpchecksum} | tail -n 1)
    else
        line=$(fgrep '.img.bz2' ${tmpchecksum} | \
               fgrep -v 'gui-' | tail -n 1)
    fi
    latest_filename=$(echo "${line}" | awk -F'[()]' '{ print $2 }')
    latest_md5=$(echo "${line}" | awk '{ print $4 }')
    rm -f ${tmpchecksum}
    debug "_filename='${latest_filename}'; _md5='${latest_md5}'"
    echo "_filename='${latest_filename}'; _md5='${latest_md5}'"
}

# Extract the version from image filename
get_version_filename() {
    [ $# -eq 2 ] || \
        error ${EC_ARGS} "get_version_filename: invalid arguments: $@"
    local branch="$1"
    local filename="$2"
    local version
    if is_master_branch "${branch}"; then
        version=$(echo "${filename}" | cut -d'-' -f5 | cut -d'.' -f1-5)
    else
        version=$(echo "${filename}" | cut -d'-' -f3 | cut -d'_' -f1)
        version=${version#v}
    fi
    echo ${version}
}

# Compare between two version strings
# Parameters: ver1 ver2
# Returns values:
# * 0 : ver1 = ver2
# * 1 : ver1 < ver2
# * 2 : ver1 > ver2
compare_version() {
    [ $# -eq 2 ] || \
        error ${EC_ARGS} "compare_version: invalid arguments: $@"
    local ver1="$1"
    local ver2="$2"
    local ver_low=$(echo -e "${ver1}\n${ver2}" | sort -V | head -n 1)
    if [ "${ver1}" = "${ver2}" ]; then
        echo 0
    elif [ "${ver1}" = "${ver_low}" ]; then
        echo 1
    else
        echo 2
    fi
}

# Checksum the image file
#
# checksum_image(file, md5)
#
# Returns:
# * 0 : file exists and its md5 hash matches the given one.
# * 1 : otherwise
checksum_image() {
    [ $# -eq 2 ] || \
        error ${EC_ARGS} "checksum_image: invalid arguments: $@"
    local file="$1"
    local md5_match="$2"
    local md5
    [ -f "${file}" ] || \
        error ${EC_NOFILE} "checksum_image: file not exists: ${file}"
    md5=$(md5 -q "${file}")
    if [ "${md5}" = "${md5_match}" ]; then
        return 0
    else
        return 1
    fi
}

# Download the latest system image (IMG file)
#
# download_image(url, outfile)
#
download_image() {
    [ $# -eq 2 ] || \
        error ${EC_ARGS} "download_image: invalid arguments: $@"
    local url="$1"
    local outfile="$2"
    local outdir=$(dirname "${outfile}")
    [ ! -d "${outdir}" ] && mkdir "${outdir}"
    echo "Downloading the new system image ..."
    echo "  <= ${url}"
    echo "  => ${outfile}"
    fetch -o "${outfile}" "${url}" \
        && echo "DONE" \
        || exit ${EC_FETCH}
}

# Mount the downloaded image (IMG file)
#
# mount_image(imgfile, mntpnt)
#
mount_image() {
    [ $# -eq 2 ] || \
        error ${EC_ARGS} "mount_image: invalid arguments: $@"
    local imgfile="$1"
    local mntpnt="$2"
    local vn=$(vnconfig -l | fgrep "not in use" | head -n 1 | cut -d':' -f 1)
    [ ! -d "${mntpnt}" ] && mkdir "${mntpnt}"
    echo "Mounting image ${imgfile} to ${mntpnt} ..."
    vnconfig -v -c ${vn} ${imgfile} || exit ${EC_VN}
    mount -r /dev/${vn}s2a ${mntpnt} \
        && echo "DONE" \
        || exit ${EC_MOUNT}
}

# Get the vn device name of the mounted image
get_vn_devname() {
    [ $# -eq 1 ] || \
        error ${EC_ARGS} "get_vn_devname: invalid arguments: $@"
    local mntpnt="$1"
    local dev=$(mount | fgrep "${mntpnt}" | cut -d' ' -f 1 | cut -d'/' -f 3)
    echo ${dev%s??}
}

# Get the filename configured for the vn device
get_vn_filename() {
    [ $# -eq 1 ] || \
        error ${EC_ARGS} "get_vn_filename: invalid arguments: $@"
    local vn="$1"
    echo $(vnconfig -l ${vn} | awk '{ print $3 }')
}

# Umount the image
#
# umount_image(mntpnt)
#
umount_image() {
    [ $# -eq 1 ] || \
        error ${EC_ARGS} "umount_image: invalid arguments: $@"
    local mntpnt="$1"
    local vn=$(get_vn_devname ${mntpnt})
    echo -n "Umounting image from ${mntpnt} ... "
    umount ${mntpnt} && echo "DONE" || exit ${EC_UMOUNT}
    echo "Disable and unconfigure VN device ${vn} ... "
    vnconfig -v -u ${vn} \
        && echo "DONE" \
        || exit ${EC_VN}
}

# Backup the old kernel
backup_kernel() {
    [ $# -eq 0 ] || \
        error ${EC_ARGS} "backup_kernel: invalid arguments: $@"
    local kerndir="/boot/kernel"
    local oldkerndir="${kerndir}.old"
    echo "Backing up current kernel to ${oldkerndir} ..."
    if [ -d "${oldkerndir}" ]; then
        warn "Previously backed kernel already exists!"
        rm -r ${oldkerndir}
        warn "Removed previously backed kernel: ${oldkerndir}"
    fi
    mkdir -p ${oldkerndir}
    chflags noschg ${kerndir}/kernel
    objcopy --strip-debug ${kerndir}/kernel ${oldkerndir}/kernel
    for f in ${kerndir}/*.ko; do
        objcopy --strip-debug ${f} ${oldkerndir}/${f##*/}
    done
    [ -f "${kerndir}/initrd.img" ] && \
        cp -p ${kerndir}/initrd.img ${oldkerndir}
    [ -f "${kerndir}/initrd.img.gz" ] && \
        cp -p ${kerndir}/initrd.img.gz ${oldkerndir}
    echo "DONE"
}

# Backup the old world
#
# backup_world(backfile)
#
backup_world() {
    [ $# -eq 1 ] || \
        error ${EC_ARGS} "backup_world: invalid arguments: $@"
    local backfile="$1"
    local backdir=$(dirname "${backfile}")
    echo "Backing up current world to ${backfile} ..."
    [ ! -d "${backdir}" ] && mkdir ${backdir}
    if [ -f "${backfile}" ]; then
        warn "Previously backed world exists!"
        rm -f "${backfile}"
        warn "Removed previously backed world: ${backfile}"
    fi
    tar -czf "${backfile}" \
        --options gzip:compression-level=1 \
        -C / \
        etc \
        bin sbin lib libexec \
        usr/bin usr/sbin usr/lib usr/libexec \
        && echo "DONE" \
        || exit ${EC_TAR}
}

# Install the new system (kernel and world, exclude /etc configurations)
install_system() {
    [ $# -eq 0 ] || \
        error ${EC_ARGS} "install_system: invalid arguments: $@"
    local file file2 item
    echo "Installing the new kernel and world ..."
    echo "  => Backing up the files to protect from overriding ..."
    for file in ${FILES_IGNORE}; do
        if [ -f "${file}" ]; then
            file2="${file}.${BAK_SUF}"
            cp -pf "${file}" "${file2}"  # NOTE: do NOT use "mv"
            echo "   * ${file}  <backed>"
        fi
    done

    echo "  => Creating distribution directories ..."
    mtree -deUq -f ${MNT_DIR}/etc.hdd/mtree/BSD.root.dist \
          -p /            || exit ${EC_MTREE}
    mtree -deUq -f ${MNT_DIR}/etc.hdd/mtree/BSD.var.dist \
          -p /var         || exit ${EC_MTREE}
    mtree -deUq -f ${MNT_DIR}/etc.hdd/mtree/BSD.usr.dist \
          -p /usr         || exit ${EC_MTREE}
    mtree -deUq -f ${MNT_DIR}/etc.hdd/mtree/BSD.include.dist \
          -p /usr/include || exit ${EC_MTREE}

    echo "  => Installing kernel and world ..."
    for item in \
            COPYRIGHT \
            bin \
            boot \
            compat \
            lib \
            libexec \
            rescue \
            sbin \
            usr/Makefile \
            usr/bin \
            usr/games \
            usr/include \
            usr/lib \
            usr/libdata \
            usr/libexec \
            usr/sbin \
            usr/share \
            var/msgs \
            var/yp; do
        echo -n "   * Installing: ${item} ... "
        cpdup -o -u ${MNT_DIR}/${item} /${item} || exit ${EC_CPDUP}
        echo "ok"
    done

    echo "  => Recovering the backed files ..."
    for file in ${FILES_IGNORE}; do
        file2="${file}.${BAK_SUF}"
        if [ -f "${file2}" ]; then
            mv -f "${file2}" "${file}"
            echo "   * ${file}  <recovered>"
        fi
    done
    echo "  => DONE!"
}

# Upgrade the system with new configuration files
upgrade_system() {
    [ $# -eq 0 ] || \
        error ${EC_ARGS} "upgrade_system: invalid arguments: $@"
    local etcdir="${CACHE_DIR}/etc.new"
    local file file_etc file_new
    [ ! -d "${CACHE_DIR}" ] && mkdir "${CACHE_DIR}"
    echo "Upgrading system ..."
    echo "  => Coping new /etc to: ${etcdir}"
    cpdup -o -u ${MNT_DIR}/etc.hdd ${etcdir} || exit ${EC_CPDUP}

    echo "  => Removing ignored files ..."
    for file_etc in ${FILES_IGNORE}; do
        file_new="${etcdir}/${file_etc#/etc/}"
        if [ -f "${file_new}" ]; then
            rm -f "${file_new}"
            echo "   * ${file_new}  <ignored>"
        fi
    done
    echo "  => Identifying new/updated config files ..."
    (cd "${etcdir}" && find -s . -type f) | while read -r file; do
        file_etc="/etc/${file#./}"
        file_new="${etcdir}/${file#./}"
        if [ -f "${file_etc}" ]; then
            if cmp -s "${file_etc}" "${file_new}"; then
                rm -f "${file_new}"
            else
                mv "${file_new}" "${file_new}.${NEW_SUF}"
                echo "   * ${file_new}  [UPDATED]"
            fi
        else
            echo "   * ${file_new}  [NEW]"
        fi
    done

    echo "  => Installing new configurations ..."
    cpdup -o -u ${etcdir} /etc || exit ${EC_CPDUP}
    echo "  => DONE!"
    rm -rf "${etcdir}"
    echo "+---------+"
    echo "| WARNING | Files with '${NEW_SUF}' suffix need manual merge!"
    echo "+---------+"
}

# Clean up obsolete and deprecated files
cleanup() {
    [ $# -eq 0 ] || \
        error ${EC_ARGS} "cleanup: invalid arguments: $@"
    local mk_upgrade tmpfile item itemcat
    mk_upgrade=/etc/upgrade/Makefile_upgrade.inc
    [ -e "${mk_upgrade}.${NEW_SUF}" ] && mk_upgrade=${mk_upgrade}.${NEW_SUF}
    tmpfile=$(mktemp -t ${NAME}) || exit ${EC_TMPFILE}
    echo "Removing obsolete and deprecated files ..."
    echo "(according to ${mk_upgrade})"
    make -f ${mk_upgrade} -V TO_REMOVE      | tr ' ' '\n' >  ${tmpfile}
    make -f ${mk_upgrade} -V TO_REMOVE_LATE | tr ' ' '\n' >> ${tmpfile}

    # Credit: https://stackoverflow.com/a/10929511
    # [ -n "${item}" ]: do not ignore the last line if not end with a '\n'
    while IFS='' read -r item || [ -n "${item}" ]; do
        if [ -n "${item}" ] && [ -e ${item} -o -L ${item} ]; then
            echo "  * ${item}"
            chflags -Rf noschg ${item}
            rm -rf ${item}
        fi
        if echo "${item}" | grep -q '/man/man[1-9]/'; then
            itemcat=$(echo "${item}" | sed 's|/man/man|/man/cat|')
            if [ -e "${itemcat}" ]; then
                echo "  * ${itemcat}"
                chflags -Rf noschg ${itemcat}
                rm -rf ${itemcat}
            fi
        fi
    done < ${tmpfile}
    rm -f ${tmpfile}
    echo "DONE"
}

# Post-upgrade checking and report:
# * check /etc for newly installed files that need manual merge
postupgrade() {
    [ $# -eq 0 ] || \
        error ${EC_ARGS} "postupgrade: invalid arguments: $@"
    echo "+=========================================================+"
    echo "The following config files need manual merge:"
    echo "+---------------------------------------------------------+"
    find -s /etc -name "*.${NEW_SUF}" | sort

    cat << _EOF_
+---------------------------------------------------------+
After manually merge the above files, reboot into the new
system, and upgrade the packages with:
    # pkg upgrade -f
+=========================================================+
_EOF_
}


#
# Sub-command functions
#

cmd_version() {
    cat <<_EOF_
v${VERSION}
Aaron LI <aly@aaronly.me>
https://github.com/liweitianux/dfly-update
_EOF_
}

cmd_usage() {
    cat <<_EOF_
dfly-update - DragonFly BSD update tool using binary release/snapshots

Usage:
    help | --help | -h
        Show this help.
    version | --version | -v
        Show version information of this tool.
    status
        Show local installed system version and remote available version.
    download <filename> <md5>
        Download the given image and check aginst the given MD5
    mount <filepath>
        Mount the given image file
    backup
        Back up the current kernel and world
    upgrade
        Install the new kernel, world, and config files
    cleanup
        Clean up obsolete files, umount and remove image file
    fly | go
        Synthetic command to upgrade the system!
_EOF_
    echo
    cmd_version
}

cmd_status() {
    [ $# -eq 0 ] || \
        error ${EC_ARGS} "cmd_status: invalid arguments: $@"
    local branch=$(get_local_branch)
    local version=$(get_local_version)
    local branch_remote version_remote has_update ret
    if [ -z "${UPDATE_BRANCH}" ]; then
        branch_remote=${branch}
    else
        branch_remote=${UPDATE_BRANCH}
    fi
    ret="$(get_latest_image ${branch_remote})" || exit $?
    eval "${ret}"
    version_remote=$(get_version_filename ${branch_remote} ${_filename})
    cat <<_EOF_
Local installed system:
    branch: ${branch}
    version: ${version}

Remote available system:
    branch: ${branch_remote}
    version: ${version_remote}
    filename: ${_filename}
    md5: ${_md5}

_EOF_

    has_update=$(compare_version ${version} ${version_remote})
    if [ ${has_update} -eq 0 ]; then
        echo "^_^ Your DragonFly is up-to-date ^_^"
    elif [ ${has_update} -eq 1 ]; then
        echo "!!! Your DragonFly needs update !!!"
    else
        echo "??? Your DragonFly is newer than remote ???"
    fi
    echo ""

    # XXX: set the global variables for use in `cmd_fly()`
    _FILENAME="${_filename}"
    _MD5=${_md5}
    _HAS_UPDATE=${has_update}

    return ${has_update}
}

# Download the given image and check aginst the given MD5
#
# usage:
#   cmd_download <filename> <md5>
cmd_download() {
    [ $# -eq 2 ] || \
        error ${EC_ARGS} "cmd_download: invalid arguments: $@"
    local filename="$1"
    local md5="$2"
    local url=$(get_image_url ${filename})
    local filepath="${CACHE_DIR}/${filename}"
    download_image "${url}" "${filepath}"
    echo -n "MD5 checking file ... "
    checksum_image "${filepath}" "${md5}" \
        && echo "OK" \
        || error ${EC_MD5} "FAILED!"
}

# Mount the given image file
#
# usage:
#   cmd_mount <file>
cmd_mount() {
    [ $# -eq 1 ] || \
        error ${EC_ARGS} "cmd_mount: invalid arguments: $@"
    local file="$1"
    [ -f "${file}" ] || \
        error ${EC_NOFILE} "checksum_image: file not exists: ${file}"
    case "${file}" in
        *.bz2)
            echo -n "Decompressing file: ${file} ... "
            bunzip2 "${file}"
            echo "DONE"
            file="${file%.bz2}"
            ;;
    esac
    mount_image "${file}" "${MNT_DIR}"
}

# Back up the current kernel and world
cmd_backup() {
    [ $# -eq 0 ] || \
        error ${EC_ARGS} "cmd_backup: invalid arguments: $@"
    backup_kernel
    backfile="${BACK_DIR}/world.tar.gz"
    backup_world "${backfile}"
}

# Install the new kernel, world, and config files.
cmd_upgrade() {
    [ $# -eq 0 ] || \
        error ${EC_ARGS} "cmd_upgrade: invalid arguments: $@"
    install_system
    upgrade_system
}

# Clean up obsolete files, umount and remove image
cmd_cleanup() {
    [ $# -eq 0 ] || \
        error ${EC_ARGS} "cmd_cleanup: invalid arguments: $@"
    cleanup
    local vn=$(get_vn_devname ${MNT_DIR})
    local filepath=$(get_vn_filename ${vn})
    umount_image ${MNT_DIR}
    rm -f ${filepath}
    echo "Removed image file: ${filepath}"
    postupgrade
}

# Integrate all the upgrading steps -> fly :-)
cmd_fly() {
    [ $# -eq 0 ] || \
        error ${EC_ARGS} "cmd_fly: invalid arguments: $@"
    echo "Checking status ..."
    cmd_status || true
    if [ ${_HAS_UPDATE} -eq 1 ]; then
        # Need to update
        local file="${CACHE_DIR}/${_FILENAME}"
        cmd_download "${_FILENAME}" ${_MD5}
        cmd_mount ${file}
        cmd_backup
        cmd_upgrade
        cmd_cleanup
    fi
}

#
# Main
#

# Load configurations
[ -r "${CONFIGFILE}" ] && . ${CONFIGFILE}

COMMAND="$1"
case "${COMMAND}" in
    version|--version|-v)
        shift
        cmd_version
        ;;
    status)
        shift
        cmd_status
        ;;
    download)
        shift
        cmd_download "$@"
        ;;
    mount)
        shift
        cmd_mount "$@"
        ;;
    backup)
        shift
        cmd_backup
        ;;
    upgrade)
        shift
        cmd_upgrade
        ;;
    cleanup)
        shift
        cmd_cleanup
        ;;
    fly|go)
        shift
        cmd_fly
        ;;
    help|--help|-h|*)
        cmd_usage
        ;;
esac

exit 0