#!/bin/bash

# read the linux repository location
. /etc/avast/vps.conf

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

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

FORCE_FULL=0
ONLY_CHECK=0
MAX_PATCHES=100


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

do_patch_update() {
    VERSION=$(cat "$STORAGE_DIR/${PACKAGE}.lat")
    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 [ -z "${NEXT}" ]; then
            echo "Missing pointer to next patch."
            return 1
        fi

        do_download "${PACKAGE}${NEXT}.inf"
        STATUS=$?
        if [ $STATUS -ne 0 ]; then
            return $STATUS
        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}"

        do_download "${PACKAGE}_${VERSION}_${NEXT}.dif"
        STATUS=$?
        if [ $STATUS -ne 0 ]; then
            return $STATUS
        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 $STATUS
        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

        mv "${PACKAGE}.tmp" "${PACKAGE}.raw"
        STATUS=$?
        if [ $STATUS -ne 0 ]; then
            rm "${PACKAGE}.tmp"
            rm "${PACKAGE}_${VERSION}_${NEXT}.dif"
            return $STATUS
        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")
    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

    if ! mkdir -p "$AVASTDATADIR/defs/$VERSION"; then
        echo "Cannot create directory: $AVASTDATADIR/defs/$VERSION"
        return 1
    fi

    if ! tar -x -C "$AVASTDATADIR/defs/$VERSION" -f "$RAW"; then
        echo "Cannot unpack the update package."
        rm -rf "$AVASTDATADIR/defs/$VERSION"
        return 1
    fi

    TMPASWDEFS="$(dirname ${RAW})/aswdefs.ini"
    echo "[Definitions]" >"$TMPASWDEFS"
    echo "Latest=$VERSION" >>"$TMPASWDEFS"

    mv "$TMPASWDEFS" "$AVASTDATADIR/defs/aswdefs.ini"

    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"; then
        echo "Update failed."
        return 1
    fi

    LAT=$(cat "${PACKAGE}.lat")
    if [ -z "$LAT" ]; then
        echo "Cannot retrieve info about the latest version."
        return 1
    fi
    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 3
    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()
{
    DAEMONPID=$(cat "$PIDFILE" 2>/dev/null)
    if [ -n "$DAEMONPID" ] && kill -0 "$DAEMONPID" 2>/dev/null; then
        # inform the daemon about finished update
        kill -HUP "$DAEMONPID" 2>/dev/null
    else
        # remove obsolete VPS directories
        CURVERSION=$(grep -e "Latest=" $AVASTDATADIR/defs/aswdefs.ini 2>/dev/null | cut -c 8-)
        for DEFDIR in "$AVASTDATADIR"/defs/????????; do
            if [ "$DEFDIR" != "$AVASTDATADIR/defs/$CURVERSION" ]; then
                rm -rf "$DEFDIR"
            fi
        done
    fi

    return 0
}

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 "  -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 ;;
        -h | --help )           usage; exit 0 ;;
        * ) echo "Unsupported option ${1}."; usage; exit 1;;
    esac
    shift
done


# create storage directory
mkdir -p "$STORAGE_DIR"

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

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

# setup lock & temp directory cleanup
trap "rm -rf \"$TEMP_DIR\" \"$LOCK_DIR\"" EXIT

cd "$TEMP_DIR" || exit 1

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

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

exit 0
