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

# Mount the downloaded image (IMG file)
#
# mount_image(imgfile, mntpnt)
#
mount_image() {
    local imgfile="$1"
    local mntpnt="$2"
    local vn
    [ -d "${mntpnt}" ] || mkdir "${mntpnt}"
    vn=$(vnconfig -c vn ${imgfile}) || exit ${EC_VN}
    mount -r /dev/${vn}s2a ${mntpnt} || exit ${EC_MOUNT}
}

# Get the vn device name of the mounted image
get_vn_devname() {
    local mntpnt="$1"
    local dev=$(mount | fgrep "${mntpnt}" | cut -d' ' -f 1 | cut -d'/' -f 3)
    echo ${dev%s??}
}

# Umount the image
#
# umount_image(mntpnt)
#
umount_image() {
    local mntpnt="$1"
    local vn=$(get_vn_devname ${mntpnt})
    umount ${mntpnt} || exit ${EC_UMOUNT}
    echo "Disable and unconfigure VN device ${vn} ..."
    vnconfig -u ${vn} || 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}
}

# 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 files to be ignored ..."
    cpignore=$(mktemp -t ${NAME}) || exit ${EC_TMPFILE}
    for file in ${FILES_IGNORE}; do
        # NOTE: 'cpdup -X' doesn't normalize multiple '/' to be one.
        echo "${MNT_DIR%/}/${file#/}" >> ${cpignore}
        echo "   * ${file}  <ignored>"
    done

    echo "  => Installing kernel and world ..."
    for item in ${INSTALL_LIST}; do
        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

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

    echo "  => Coping new /etc to: ${etcdir}"
    ${CPDUP} -o ${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 ${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 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

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

_EOF_
}


CONFIGFILE=""
ISTART=0
ISTOP=999

while getopts :c:dhs:S: opt; do
	case ${opt} in
	c)
		CONFIGFILE="${OPTARG}"
		;;
	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))
[ $# -ne 1 ] && { usage; exit ${EC_ARGS}; }
IMGFILE="$1"
[ -r "${IMGFILE}" ] || error ${EC_NOFILE} "file not exists: ${IMGFILE}"

if [ -n "${CONFIGFILE}" ]; then
    if [ -r "${CONFIGFILE}" ]; then
        . ${CONFIGFILE}
    else
        error ${EC_CONFIG} "cannot read config file: ${CONFIGFILE}"
    fi
fi

istep=0
echo "[${istep}] Mounting image ${IMGFILE} to ${MNT_DIR} ..."
[ ${istep} -ge ${ISTART} -a ${istep} -le ${ISTOP} ] &&
    mount_image "${IMGFILE}" "${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