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

set -e

DEBUG=${DEBUG:-""}

NAME="dfly-update"
VERSION="0.?.?"
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


#
# 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__"


#
# Helper Functions
#

debug() {
    [ -n "${DEBUG}" ] && echo "DEBUG: $@" >&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
    echo "Fetch remote systems checksum: ${url_checksum}" >&2
    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: $@"
    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}" ] && md5=$(md5 -q "${file}")
    if [ -n "${md5}" ] && [ "${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 -v "${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 -v "${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 "Previous backed up kernel already exists!"
        rm -r ${oldkerndir}
        warn "Removed previously backed up 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 "Backed up old kernel at: ${oldkerndir}"
}

# 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 -v ${backdir}
    if [ -f "${backfile}" ]; then
        warn "  => Previously backed up 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 "  => Stashing the files to protect from overriding ..."
    for file in ${FILES_IGNORE}; do
        file2="${file}.${BAK_SUF}"
        echo -n "  * ${file}  "
        mv -f "${file}" "${file2}"
        echo "(stashed)"
    done

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

    echo "  => Installing kernel and world ..."
    for item in COPYRIGHT \
            bin \
            boot \
            compat \
            lib \
            libexec \
            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 "done!"
    done

    echo "  => Recovering the stashed files ..."
    for file in ${FILES_IGNORE}; do
        file2="${file}.${BAK_SUF}"
        echo -n "  * ${file}  "
        mv -f "${file2}" "${file}"
        echo "(recovered)"
    done
    echo "  => DONE"
}

}


#
# 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
_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
    if [ -z "${UPDATE_BRANCH}" ]; then
        branch_remote=${branch}
    else
        branch_remote=${UPDATE_BRANCH}
    fi
    eval "$(get_latest_image ${branch_remote})" || exit $?
    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
}

# 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 <filepath>
cmd_mount() {
    [ $# -eq 1 ] || \
        error ${EC_ARGS} "cmd_mount: invalid arguments: $@"
    local filepath="$1"
    case "${filepath}" in
        *.bz2)
            echo -n "Decompressing file: ${filepath} ..."
            bunzip2 "${filepath}"
            echo "DONE"
            filepath="${filepath%.bz2}"
            ;;
    esac
    mount_image "${filepath}" "${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}"
}


#
# 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
        ;;
    help|--help|-h|*)
        cmd_usage
        ;;
esac

exit 0