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


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

#
# Helper Functions
#

debug() {
    [ -n "${DEBUG}" ] && echo "DEBUG: $@" >&2
}

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

error() {
    echo "ERROR: $@" >&2
}

check_os() {
    local os=$(uname -s)
    if [ "${os}" != "DragonFly" ]; then
        error "Not a DragonFly BSD system!"
        exit ${EC_OS}
    fi
}

# contains(string, substring)
#
# Returns 0 if the specified string contains the specified substring,
# otherwise returns 1.
#
# Credit: https://stackoverflow.com/a/8811800
contains() {
    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() {
    if contains "$1" "-DEV-"; then
        echo "MASTER"
    else
        echo "RELEASE"
    fi
}

# Determine whether the given name refers to the master branch?
is_master_branch() {
    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() {
    check_os || exit $?
    uname -r | awk -F'-' '{ print $2 }'
}

# Get the version of local installed system
get_local_version() {
    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() {
    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() {
    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() {
    local latest_filename latest_md5 line
    local branch="$1"
    local url_checksum=$(get_md5list_url ${branch})
    local tmpchecksum=$(mktemp -t ${NAME}) || \
        exit ${EC_TMPFILE}
    echo "Fetch remote systems checksum: ${url_checksum}" >&2
    fetch -q -o ${tmpchecksum} "${url_checksum}"
    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() {
    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() {
    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() {
    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() {
    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 $?
}


#
# 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
_EOF_
    echo
    cmd_version
}

cmd_status() {
    local branch version branch_remote version_remote has_update
    branch=$(get_local_branch)
    version=$(get_local_version)
    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() {
    local filename="$1"
    local md5="$2"
    local url=$(get_image_url ${filename})
    local filepath="${CACHE_DIR}/${filename}"
    download_image "${url}" "${filepath}"
    if checksum_image "${filepath}" "${md5}"; then
        echo "Downloaded and MD5-checked image file."
    else
        echo "Downloaded image file does not match the MD5!"
        exit ${EC_MD5}
    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 "$@"
        ;;
    help|--help|-h|*)
        cmd_usage
        ;;
esac

exit 0