#!/bin/sh
#
# Copyright (c) 2017-2019 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.2.0"
NAME="dfly-update"
AUTHOR="Aaron LI <aly@aaronly.me>"
URL="https://github.com/liweitianux/dfly-update"

#
# Error Codes
#

EC_USAGE=1
EC_ARGS=2
EC_CONFIG=11
EC_TMPFILE=12
EC_MD5=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
EC_PW=22


#
# Default Configurations
#

# Path to the cpdup(1) executable
CPDUP="/bin/cpdup"
# 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}"

# List of files/directories to be installed
# NOTE: Do NOT include 'etc' below!
INSTALL_LIST='
    /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
'

# Ignored files to be kept from overriding by the upgrade.
FILES_IGNORE='
    /boot/loader.conf
    /etc/crypttab
    /etc/fstab
    /etc/group
    /etc/localtime
    /etc/login.conf.db
    /etc/master.passwd
    /etc/motd
    /etc/passwd
    /etc/pwd.db
    /etc/rc.conf
    /etc/spwd.db
    /var/db/locate.database
    /var/mail/root
'
# Filename suffix for newly installed files that need (manual) merge.
NEW_SUF="__new__"


#
# 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}
}

# Get the device mounted at the given directory
get_mount_dev() {
    local mntpnt="$1"
    local dev=$(df | awk '$NF == "'${mntpnt}'" { print $1 }')
    echo "${dev}"
}

# Mount the specified system image file
#
# mount_image(imgfile, mntpnt)
#
mount_image() {
    local imgfile="$1"
    local mntpnt="$2"
    local dev vn

    dev=$(get_mount_dev ${mntpnt})
    [ -z "${dev}" ] ||
        error ${EC_MOUNT} "${dev} already mounted at ${mntpnt}"

    [ -d "${mntpnt}" ] || mkdir "${mntpnt}"
    vn=$(vnconfig -c vn ${imgfile}) || exit ${EC_VN}
    mount -r /dev/${vn}s2a ${mntpnt} || exit ${EC_MOUNT}
}

# Umount the image and unconfigure the underlying VN device
#
# umount_image(mntpnt)
#
umount_image() {
    local mntpnt="$1"
    local dev vn
    dev=$(get_mount_dev ${mntpnt})
    vn=${dev#/dev/}
    umount ${mntpnt} || exit ${EC_UMOUNT}
    echo "Unconfigure ${vn} ..."
    vnconfig -u "${vn%s??}" || exit ${EC_VN}
}

# Backup the old kernel
backup_kernel() {
    local kerndir="/boot/kernel"
    local oldkerndir="${kerndir}.old"
    [ -d "${oldkerndir}" ] && {
        rm -r ${oldkerndir}
        warn "Removed previously backed kernel: ${oldkerndir}"
    }

    echo "Backing up current kernel to ${oldkerndir} ..."
    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() {
    local backfile="$1"
    local backdir=$(dirname "${backfile}")
    [ -d "${backdir}" ] || mkdir ${backdir}
    [ -f "${backfile}" ] && {
        rm -f "${backfile}"
        warn "Removed previously backed world: ${backfile}"
    }

    echo "Backing up current world to ${backfile} ..."
    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}
}

# Make the ignore file for cpdup(1)
make_cpignore() {
    local cpignore=$(mktemp -t ${NAME}) || exit ${EC_TMPFILE}
    local f f2
    for f in ${FILES_IGNORE}; do
        f2=$(echo "${f}" | sed 's|^/etc/|/etc.hdd/|')
        # NOTE: 'cpdup -X' doesn't normalize multiple '/' to be one.
        echo "${MNT_DIR%/}/${f2#/}" >> ${cpignore}
    done
    echo ${cpignore}
}

# Install the new system (kernel and world, excluding /etc)
install_system() {
    local file item path cpignore

    echo "  => Creating distribution directories ..."
    for item in \
            root:/ \
            var:/var \
            usr:/usr \
            include:/usr/include; do
        file=BSD.${item%:*}.dist
        path=${item#*:}
        echo "   * mtree: ${path} ..."
        mtree -deUq -f ${MNT_DIR}/etc.hdd/mtree/${file} -p ${path} ||
            exit ${EC_MTREE}
    done

    echo "  => Collecting ignored files for cpdup ..."
    cpignore=$(make_cpignore) || exit ${EC_TMPFILE}

    echo "  => Installing kernel and world ..."
    for item in ${INSTALL_LIST}; do
        [ "${item%/}" = "/etc" ] && {
            warn "'/etc' is in 'INSTALL_LIST'; force ignored"
            continue
        }
        echo -n "   * Installing: ${item} ... "
        # NOTE: 'cpdup -X' doesn't normalize multiple '/' to be one.
        ${CPDUP} -o -X ${cpignore} ${MNT_DIR%/}/${item#/} /${item} ||
            exit ${EC_CPDUP}
        echo "ok"
    done

    rm -f ${cpignore}
    echo "  => DONE!"
}

# Add new users and groups
add_users() {
    local fpasswd="${MNT_DIR}/etc.hdd/master.passwd"
    local fgroup="${MNT_DIR}/etc.hdd/group"
    local _name _pw _uid _gid _gids item
    local _class _change _expire _gecos _home _shell _members

    echo "  => Adding new users ..."
    _gids=""
    grep -Ev '^(#.*|\s*)$' ${fpasswd} |
            while IFS=':' read -r _name _pw _uid _gid _class \
                                  _change _expire _gecos _home _shell; do
        pw usershow ${_name} -q >/dev/null && continue

        # NOTE: There is circular dependence: 'groupadd' requires the members
        #       already exist, while 'useradd' requires the group exists.
        #       So first assign new users to the 'nogroup' group, and make
        #       adjustments after group creation.
        echo "   * ${_name}, ${_uid}, ${_gid}, ${_gecos}, ${_home}, ${_shell}"
        pw useradd ${_name} \
            -u ${_uid} -g nogroup -d ${_home} -s ${_shell} \
            -L "${_class}" -c "${_gecos}" || exit ${EC_PW}
        _gids="${_gids} ${_name}:${_gid}"
    done

    echo "  => Adding new groups ..."
    grep -Ev '^(#.*|\s*)$' ${fgroup} |
            while IFS=':' read -r _name _pw _gid _members; do
        pw groupshow ${_name} -q >/dev/null && continue

        echo "   * ${_name}, ${_gid}, ${_members}"
        pw groupadd ${_name} -g ${_gid} -M "${_members}" || exit ${EC_PW}
    done

    echo "  => Adjusting the group of new users ..."
    for item in ${_gids}; do
        _name=${item%:*}
        _gid=${item#*:}
        echo "   * ${_name}, ${_gid}"
        pw usermod ${_name} -g ${_gid} || exit ${EC_PW}
    done
}

# Upgrade the system with new configuration files
upgrade_system() {
    local etcdir file file_etc file_new cpignore

    [ -d "${CACHE_DIR}" ] || mkdir "${CACHE_DIR}"
    etcdir="${CACHE_DIR}/etc.new"
    rm -rf "${etcdir}"

    echo "  => Collecting ignored files for cpdup ..."
    cpignore=$(make_cpignore) || exit ${EC_TMPFILE}
    echo "  => Coping new /etc to: ${etcdir}"
    ${CPDUP} -o -X ${cpignore} ${MNT_DIR%/}/etc.hdd ${etcdir} ||
        exit ${EC_CPDUP}
    rm -f ${cpignore}

    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 ${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() {
    local mk_upgrade tmpfile item itemcat

    mk_upgrade=/etc/upgrade/Makefile_upgrade.inc
    [ -e "${mk_upgrade}.${NEW_SUF}" ] && mk_upgrade=${mk_upgrade}.${NEW_SUF}
    echo "Removing obsolete and deprecated files ..."
    echo "(according to ${mk_upgrade})"
    tmpfile=$(mktemp -t ${NAME}) || exit ${EC_TMPFILE}
    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"
}

# Misc operations after upgrade
postupgrade() {
    echo "Rebuild capability database ..."
    cap_mkdb /etc/login.conf
    echo "Rebuild password database ..."
    pwd_mkdb -p /etc/master.passwd
    echo "Rebuild mail aliases db ..."
    newaliases
    echo "Rebuild whatis database ..."
    makewhatis
    echo "Rebuild shared library cache ..."
    ldconfig -R

    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 update
    # pkg upgrade [-f]
+=========================================================+
_EOF_
}


#
# Main
#

usage() {
    cat <<_EOF_
Upgrade a DragonFly BSD system using a binary release/snapshot

Usage:
${0##*/} [-h] [-c <config-file>] [-d]
    [-s <start-step>] [-S <stop-step] {dfly.img}

Options:
    -h : show this help
    -c : load the specified configuration file
    -d : enable the debug mode
    -s : start at the specified step
    -S : stop at the specified step (inclusive)

Steps:
0. mount the system image
1. backup current kernel
2. backup current world
3. install new kernel and world (excluding /etc)
4. add new users and groups
5. install and update /etc files
6. clean up obsolete files
7. umount the system image
8. misc operations after upgrade

Note:
The {dfly.img} is required when <start-step> is 0 (the default),
otherwise it is assumed that the image file is already mounted
at \${MNT_DIR} (default to be '${MNT_DIR}').

v${VERSION}
${AUTHOR}
${URL}

_EOF_
}


ISTART=0
ISTOP=99

while getopts :c:dhs:S: opt; do
    case ${opt} in
    c)
        CONFIGFILE="${OPTARG}"
        if [ -r "${CONFIGFILE}" ]; then
            . ${CONFIGFILE}
        else
            error ${EC_CONFIG} "cannot read config file: ${CONFIGFILE}"
        fi
        ;;
    d)
        DEBUG=yes
        ;;
    h)
        usage
        exit ${EC_USAGE}
        ;;
    s)
        ISTART="${OPTARG}"
        ;;
    S)
        ISTOP="${OPTARG}"
        ;;
    \?)
        log "Invalid option -${OPTARG}"
        usage
        exit ${EC_ARGS}
        ;;
    :)
        log "Option -${OPTARG} requires an argument"
        usage
        exit ${EC_ARGS}
        ;;
    esac
done

shift $((OPTIND - 1))

if [ ${ISTART} -eq 0 ]; then
    [ $# -eq 1 ] || error ${EC_ARGS} "no image file provided"
    IMG="$1"
    [ -f "${IMG}" ] || error ${EC_NOFILE} "image file not exists: ${IMG}"
else
    DEV=$(get_mount_dev ${MNT_DIR})
    [ -n "${DEV}" ] || error ${EC_MOUNT} "nothing mounted at: ${MNT_DIR}"
fi

istep=0
echo "[${istep}] Mounting image to ${MNT_DIR} ..."
[ ${istep} -ge ${ISTART} -a ${istep} -le ${ISTOP} ] &&
    mount_image "${IMG}" "${MNT_DIR}" ||
    echo "(skipped)"

istep=$((${istep} + 1))
echo "[${istep}] Backing up current kernel ..."
[ ${istep} -ge ${ISTART} -a ${istep} -le ${ISTOP} ] &&
    backup_kernel ||
    echo "(skipped)"

istep=$((${istep} + 1))
echo "[${istep}] Backing up current world ..."
[ ${istep} -ge ${ISTART} -a ${istep} -le ${ISTOP} ] &&
    backup_world "${BACK_DIR}/world.tar.gz" ||
    echo "(skipped)"

istep=$((${istep} + 1))
echo "[${istep}] Installing new kernel and world ..."
[ ${istep} -ge ${ISTART} -a ${istep} -le ${ISTOP} ] &&
    install_system ||
    echo "(skipped)"

istep=$((${istep} + 1))
echo "[${istep}] Adding new users and groups ..."
[ ${istep} -ge ${ISTART} -a ${istep} -le ${ISTOP} ] &&
    add_users ||
    echo "(skipped)"

istep=$((${istep} + 1))
echo "[${istep}] Upgrade system files ..."
[ ${istep} -ge ${ISTART} -a ${istep} -le ${ISTOP} ] &&
    upgrade_system ||
    echo "(skipped)"

istep=$((${istep} + 1))
echo "[${istep}] Clean up obsolete files ..."
[ ${istep} -ge ${ISTART} -a ${istep} -le ${ISTOP} ] &&
    cleanup ||
    echo "(skipped)"

istep=$((${istep} + 1))
echo "[${istep}] Umounting image from ${MNT_DIR} ..."
[ ${istep} -ge ${ISTART} -a ${istep} -le ${ISTOP} ] &&
    umount_image "${MNT_DIR}" ||
    echo "(skipped)"

istep=$((${istep} + 1))
echo "[${istep}] Misc operations after upgrade ..."
[ ${istep} -ge ${ISTART} -a ${istep} -le ${ISTOP} ] &&
    postupgrade ||
    echo "(skipped)"

exit 0