#!/bin/bash

# read the VPS repository URL
. /etc/avast/vps.conf

PACKAGE="vps9"
REPO_URL="$URL/$PACKAGE"

AVASTDATADIR="/var/lib/avast"
STORAGE_DIR="$AVASTDATADIR/Setup/filedir"
DAEMON="/usr/bin/avast"
BSPATCH="/usr/lib/avast/bspatch"
SUBMIT="/usr/lib/avast/submit"
MD5SUM="md5sum"
DOWNLOAD="curl -L -s --retry 10 --connect-timeout 30 -f"
DOWNLOAD_RETRIES=2
PIDFILE="/run/avast/avast.pid"
LOCK_DIR="$AVASTDATADIR/Setup/lock"
TEMP_DIR="$AVASTDATADIR/Setup/tmp"

FORCE_FULL=0
ONLY_CHECK=0
MAX_PATCHES=100
EXTENDED_EXIT_CODES=0


do_download() {
    local i
    local FN="$1"
    for ((i=0; i<=DOWNLOAD_RETRIES; i++)); do
        if [ "$i" -ne 0 ]; then
            echo "Downloading $FN... (retry $i)"
        fi
        $DOWNLOAD "$REPO_URL/$FN" -o "$FN"
        STATUS=$?
        if [[ $STATUS -eq 0 ]]; then
            if [ -f "$FN" ]; then
                return 0
            fi
            STATUS=1
        elif [[ $STATUS -eq 22 ]]; then
            # Do not retry on HTTP error
            break
        fi
    done
    echo "Download failed: $DOWNLOAD $REPO_URL/$FN [exit code $STATUS]"
    return "$STATUS"
}

validate_vps_version() {
    local VER="$1"
    if [[ -z "$VER" ]]; then
        echo "Cannot retrieve info about the latest version."
        return 1
    fi
    if ! [[ "$VER" =~ ^[0-9]+$ ]]; then
        echo "Invalid VPS version: $VER"
        return 1
    fi
    return 0
}

do_patch_update() {
    VERSION=$(cat "$STORAGE_DIR/${PACKAGE}.lat")
    local RAW="$STORAGE_DIR/${PACKAGE}.raw"

    if [[ -z "${VERSION}"  ||  ! -f "${RAW}" ]]; then
        echo "Starting raw package not present."
        return 1
    fi

    if ! do_download "${PACKAGE}${VERSION}.inf"; then
        return 1
    fi

    CURRENT=${VERSION}
    SEQUENCE=()
    while true; do
        NEXT=$(grep "NEXT=" "${PACKAGE}${CURRENT}.inf" | tail -n 1 | cut -d= -f2)
        if ! validate_vps_version "${NEXT}" ; then
            echo "Missing pointer to next patch."
            return 1
        fi

        if ! do_download "${PACKAGE}${NEXT}.inf" ; then
            return 1
        fi

        SEQUENCE+=("${NEXT}")
        CURRENT="${NEXT}"

        if [ "${CURRENT}" = "${LAT}" ]; then
            break
        fi

        if [ "${#SEQUENCE[@]}" -gt "${MAX_PATCHES}" ]; then
            echo "Too many patches. Falling back to full update."
            return 1
        fi
    done

    echo "Will apply ${#SEQUENCE[@]} patches..."

    for NEXT in "${SEQUENCE[@]}"; do
        echo "Applying patch: ${VERSION} -> ${NEXT}"

        if ! do_download "${PACKAGE}_${VERSION}_${NEXT}.dif"; then
            return 1
        fi

        "$BSPATCH" "$RAW" "${PACKAGE}.tmp" "${PACKAGE}_${VERSION}_${NEXT}.dif" 2>/dev/null
        STATUS=$?
        if [ $STATUS -ne 0 ]; then
            echo "bspatch failed ($STATUS)"
            rm -f "${PACKAGE}.tmp"
            rm "${PACKAGE}_${VERSION}_${NEXT}.dif"
            return 1
        fi

        MD5=$(grep "MD5=" "${PACKAGE}${NEXT}.inf" | tail -n 1 | cut -d= -f2)
        if [ "$MD5" != "$("${MD5SUM}" "${PACKAGE}.tmp" | cut -d ' ' -f1)" ]; then
            echo "MD5 checksum does not match."
            rm "${PACKAGE}.tmp"
            rm "${PACKAGE}_${VERSION}_${NEXT}.dif"
            return 1
        fi

        if ! mv "${PACKAGE}.tmp" "${PACKAGE}.raw" ; then
            rm "${PACKAGE}.tmp"
            rm "${PACKAGE}_${VERSION}_${NEXT}.dif"
            return 1
        fi

        RAW="${PACKAGE}.raw"
        VERSION="${NEXT}"

        if [ "${VERSION}" = "${LAT}" ]; then
            return 0
        fi
    done
}

do_full_update() {
    echo "Updating to: ${LAT}"
    if ! do_download "${PACKAGE}${LAT}.inf"; then
        return 1
    fi
    MD5=$(grep "MD5=" "${PACKAGE}${LAT}.inf" | tail -n 1 | cut -d= -f2)
    if [ -z "$MD5" ]; then
        return 1
    fi

    if ! do_download "${PACKAGE}${LAT}.ful"; then
        return 1
    fi

    if ! gunzip -c "${PACKAGE}${LAT}.ful" > "${PACKAGE}.tmp"; then
        return 1
    fi

    if [ "$MD5" != "$(${MD5SUM} ${PACKAGE}.tmp | cut -d ' ' -f1)" ]; then
        echo "MD5 checksum does not match."
        return 1
    fi

    if ! mv "${PACKAGE}.tmp" "${PACKAGE}.raw"; then
        return 1
    fi

    VERSION="${LAT}"
    return 0
}

do_update_vps() {
    VERSION=$(cat "$STORAGE_DIR/${PACKAGE}.lat")
    local RAW="$STORAGE_DIR/${PACKAGE}.raw"

    if [[ -z "${VERSION}"  ||  ! -f "${RAW}" ]]; then
        echo "Initial raw package not present."
        return 1
    fi

    if [ -d "$AVASTDATADIR/defs/$VERSION" ]; then
        echo "Directory already exists: $AVASTDATADIR/defs/$VERSION"
        return 1
    fi

    local EXTRACT_DIR="$TEMP_DIR/defs/$VERSION"
    rm -rf "$EXTRACT_DIR"
    mkdir -p "$EXTRACT_DIR"
    if ! tar -x -C "$EXTRACT_DIR" -f "$RAW"; then
        echo "Cannot unpack the update package."
        return 1
    fi

    printf "[Definitions]\nLatest=%s\n" "$VERSION" >"$TEMP_DIR/defs/aswdefs.ini"

    # verify VPS authenticity and integrity
    if ! "$DAEMON" -d "$TEMP_DIR"; then
        echo "Error: Corrupted VPS."
        return 1
    fi

    # deploy new VPS
    if ! mkdir -p "$AVASTDATADIR/defs"; then
        echo "Cannot create directory: $AVASTDATADIR/defs"
        return 1
    fi
    mv "$TEMP_DIR/defs/$VERSION" "$AVASTDATADIR/defs/"
    mv "$TEMP_DIR/defs/aswdefs.ini" "$AVASTDATADIR/defs/"

    return 0
}

update() {
    # Download the latest VPS package
    echo "Connecting to repository: $REPO_URL"

    if [ -f "$STORAGE_DIR/${PACKAGE}.lat" ]; then
        VERSION=$(cat "$STORAGE_DIR/${PACKAGE}.lat")
    fi
    echo "Current VPS version: $VERSION"

    if ! do_download "${PACKAGE}.lat" || ! validate_vps_version "$(cat "${PACKAGE}.lat")"; then
        echo "Update failed."
        return 1
    fi

    LAT=$(cat "${PACKAGE}.lat")
    if [ "$LAT" = "$VERSION" ]; then
        echo "VPS is up to date."
        return 2
    else
        echo "Latest VPS version:  $LAT"
    fi

    if [ "${ONLY_CHECK}" -eq 1 ]; then
        return 0
    fi

    STATUS=1
    # Do incremental update, if we have all prerequisites:
    if [[ "${FORCE_FULL}" -eq 0  &&  -f "$STORAGE_DIR/$PACKAGE.raw"  &&  -f "$STORAGE_DIR/$PACKAGE.lat" ]]; then
        echo "Trying incremental updates..."
        do_patch_update
        STATUS=$?
    fi

    # Fall back to full update:
    if [ $STATUS -ne 0 ]; then
        if [ "${FORCE_FULL}" -eq 0 ]; then
            echo "Trying full update..."
        fi
        do_full_update
        STATUS=$?
        if [ $STATUS -ne 0 ]; then
            echo "Update failed."
            return 1
        fi
    fi

    # Update the storage (or render it incomplete on a failure):
    rm -f "$STORAGE_DIR/${PACKAGE}.raw" "$STORAGE_DIR/${PACKAGE}.lat" 2>/dev/null
    mv -f "${PACKAGE}.raw" "${PACKAGE}.lat" "$STORAGE_DIR" 2>/dev/null
    STATUS=$?
    if [ $STATUS -ne 0 ]; then
        echo "Cannot replace ${PACKAGE}.raw and ${PACKAGE}.lat files"
        return 1
    fi

    return 0
}

restart_avast()
{
    local DAEMONPID
    DAEMONPID=$(cat "$PIDFILE" 2>/dev/null)
    if [[ -z "${DAEMONPID}" ]]; then
        return 0
    fi

    # This fails if we're running as avast user, but daemon is running as root
    if kill -0 "${DAEMONPID}" 2>/dev/null; then
        kill -HUP "${DAEMONPID}" 2>/dev/null
        return 0
    fi

    if [[ ${EXTENDED_EXIT_CODES} -eq 1 ]]; then
        # Message disabled by --exit-code (so we don't get it after dropping root privileges)
        return 0
    fi

    echo
    echo "Please reload avast:"
    echo "  systemctl reload avast.service"
    echo "or"
    echo "  kill -HUP ${DAEMONPID}"
    return 0
}

delete_old_vps()
{
    local VPS_DIR="${AVASTDATADIR}/defs"
    if [[ -d "$VPS_DIR" ]]; then
        # shellcheck disable=SC2010
        ls -r "$VPS_DIR" | grep -x '[0-9]\{8\}' | tail -n +5 | \
        while read -r VER; do
            if [ -d "${VPS_DIR}/${VER}" ]; then
                echo "Removing old VPS: ${VER}"
                rm -rf "${VPS_DIR:?}/${VER}"
            fi
        done
    fi
}

usage()
{
    echo "vpsupdate.sh [--check | --full]"
    echo "Options:"
    echo "  -c, --check         Only check for new VPS, do not update"
    echo "  -f, --full          Force full update (instead of incremental)"
    echo "  -p, --max-patches   Limit incremental updates (default limit is $MAX_PATCHES)"
    echo "  -x, --exit-code     Enable advanced exit codes: 0: updated / update available (--check), 1: failed, 2: up-to-date"
    echo "  -h, --help          Show this help"
}

# Parse options:
while [ "$1" != "" ]; do
    case "$1" in
        -c | --check )          ONLY_CHECK=1 ;;
        -f | --full )           FORCE_FULL=1 ;;
        -p | --max-patches )    MAX_PATCHES=$2; shift ;;
        -x | --exit-code )      EXTENDED_EXIT_CODES=1 ;;
        -h | --help )           usage; exit 0 ;;
        * ) echo "Unsupported option ${1}."; usage; exit 1;;
    esac
    shift
done

# drop root privileges (if we have `avast` user)
if [[ $UID -eq 0 ]] && id avast >/dev/null 2>&1 ; then
    runuser -u avast -- "$0" "$@" --exit-code
    STATUS=$?
    if [[ ${STATUS} -eq 0 && ${ONLY_CHECK} -eq 0 ]]; then  # updated
        restart_avast
        exit 0
    elif [[ ${STATUS} -eq 1 ]]; then  # failed
        exit 1
    else  # up-to-date, or --check: update available
        exit 0
    fi
fi

# create storage directory
mkdir -p "$STORAGE_DIR"

# clean up temp dirs on exit
trap "rm -rf \"$DOWNLOAD_DIR\" \"$TEMP_DIR\" \"$LOCK_DIR\"" EXIT

# create lock directory to prevent parallel run
if ! mkdir "$LOCK_DIR" 2>/dev/null; then
    exit 1
fi

# create temporary work directory
DOWNLOAD_DIR="$(mktemp -d --tmpdir=/tmp -q -t avast.download-XXXXX)"

cd "$DOWNLOAD_DIR" || exit 1

# perform update
if update; then
    if [[ ${ONLY_CHECK} -eq 1 ]]; then
        exit 0
    fi
    # New VPS available, update to it
    if do_update_vps; then
        echo "Update successful."
        delete_old_vps
        restart_avast
    fi
else
    STATUS=$?
    if [[ ${EXTENDED_EXIT_CODES} -eq 1 ]]; then
        exit $STATUS
    fi
    if [ $STATUS -eq 1 ]; then
        # Download/package error
        exit 1
    elif [ $STATUS -eq 2 ]; then
        # VPS up to date
        :
    fi
fi

# send submits if any
$SUBMIT --flush --quiet

exit 0
