#!/usr/bin/bash

set -o pipefail
set -o errexit

# Script to upgrade ImunifyAV and/or Imunify360
# bypassing gradual rollout

VERSION="1.4"
PACKAGES=''

PRODUCT_AV='imunify-antivirus'
PRODUCT_I360='imunify360-firewall'

IMUNIFY_ANTIVIRUS_PACKAGES="imunify-antivirus \
    ai-bolit \
    alt-common-release \
    alt-php-hyperscan \
    imunify-release \
    imunify-common \
    imunify-notifier \
    imunify-core \
    imunify-realtime-av \
    imunify-realtime-av-imrt2 \
    rustbolit \
    imunify-ui \
    imunify-wp-security \
    imunify360-venv \
    alt-php-internal \
    app-version-detector"

IMUNIFY360_FIREWALL_PACKAGES="imunify360-firewall \
    imunify-antivirus \
    ai-bolit \
    alt-common-release \
    alt-php-hyperscan \
    imunify-release \
    imunify-common \
    imunify-notifier \
    imunify-core \
    imunify-realtime-av \
    imunify-realtime-av-imrt2 \
    rustbolit \
    imunify-ui \
    imunify-wp-security \
    imunify360-venv \
    cloudlinux-backup-utils \
    imunify360-ossec \
    imunify360-pam \
    imunify360-php-i360 \
    imunify360-webshield-bundle \
    imunify360-unified-access-logger \
    alt-php-internal \
    app-version-detector"

IMUNIFY360_FIREWALL_RPM_ONLY_PACKAGES="imunify-auditd-log-reader \
    minidaemon"

IMUNIFY_ANTIVIRUS_RPM_ONLY_PACKAGES="minidaemon"

REPORT_CMD="${REPORT_CMD:-/usr/libexec/report-command-error}"
DEBIAN_VERSION_FILE="${DEBIAN_VERSION_FILE:-/etc/debian_version}"
if [ -f "$DEBIAN_VERSION_FILE" ]; then
    INTEGRITY_SCRIPT_PATH="${INTEGRITY_SCRIPT_PATH:-/opt/imunify360/venv/share/imunify360/scripts/imunify-check-pkg-integrity}"
else
    INTEGRITY_SCRIPT_PATH="${INTEGRITY_SCRIPT_PATH:-/usr/share/imunify360/scripts/imunify-check-pkg-integrity}"
fi

# Source the integrity script for shared functions:
# filter_verify_noise, filter_whitelisted_files, check_dpkg_verify_supported,
# INTEGRITY_WHITELIST_FILES
if [[ -f "$INTEGRITY_SCRIPT_PATH" ]]; then
    source "$INTEGRITY_SCRIPT_PATH"
fi

verify_single_package() {
    local pkg="$1"
    local raw_output=""
    case "$OSTYPE" in
        centos)
            raw_output=$(rpm -V "$pkg" 2>&1) || raw_output="${raw_output}"
            ;;
        debian)
            raw_output=$(dpkg -V "$pkg" 2>&1) || raw_output="${raw_output}"
            ;;
    esac
    local filtered=""
    filtered=$(echo "$raw_output" | filter_verify_noise | filter_whitelisted_files)
    if [[ -n "$filtered" ]]; then
        echo "$filtered"
        return 1
    fi
    return 0
}

reinstall_package_centos() {
    local pkg="$1"
    echo "  Reinstalling $pkg via yum..."
    if ! yum -y reinstall "$pkg" --enablerepo="imunify360-rollout-*-bypass"; then
        echo "  WARNING: yum reinstall $pkg failed"
        return 1
    fi
    return 0
}

reinstall_package_debian() {
    local pkg="$1"
    echo "  Reinstalling $pkg via apt-get..."
    if ! apt-get install --reinstall -y "$pkg"; then
        echo "  WARNING: apt-get reinstall $pkg failed"
        return 1
    fi
    return 0
}

check_and_reinstall_broken_packages() {
    if ! type filter_verify_noise &>/dev/null; then
        echo "Integrity check: integrity script not available, skipping"
        return 0
    fi

    if [[ -z "${PACKAGES:-}" ]]; then
        return 0
    fi

    if [[ "$OSTYPE" == "debian" ]]; then
        if ! check_dpkg_verify_supported; then
            echo "Integrity check: dpkg -V not supported, skipping"
            return 0
        fi
    fi

    echo "Integrity check: verifying packages after update..."
    local had_broken=false

    for pkg in $PACKAGES; do
        local verify_output=""
        local verify_rc=0
        verify_output=$(verify_single_package "$pkg" 2>&1) || verify_rc=$?

        if [[ $verify_rc -ne 0 ]] && [[ -n "$verify_output" ]]; then
            if echo "$verify_output" | grep -q "is not installed"; then
                continue
            fi
            had_broken=true
            echo "  BROKEN: $pkg"
            echo "$verify_output" | sed 's/^/    /'

            local reinstall_rc=0
            case "$OSTYPE" in
                centos) reinstall_package_centos "$pkg" || reinstall_rc=$? ;;
                debian) reinstall_package_debian "$pkg" || reinstall_rc=$? ;;
            esac

            if [[ $reinstall_rc -eq 0 ]]; then
                local recheck_output=""
                local recheck_rc=0
                recheck_output=$(verify_single_package "$pkg" 2>&1) || recheck_rc=$?
                if [[ $recheck_rc -eq 0 ]]; then
                    echo "  Fixed: $pkg"
                else
                    echo "  WARNING: $pkg could not fix after reinstall, still broken"
                fi
            fi
        fi
    done

    if [[ "$had_broken" == "true" ]]; then
        echo "Integrity check: broken packages found, running Sentry report..."
        if [[ -x "${REPORT_CMD}" ]] && [[ -x "${INTEGRITY_SCRIPT_PATH}" ]]; then
            "${REPORT_CMD}" "${INTEGRITY_SCRIPT_PATH}" || true
        fi
    else
        echo "Integrity check: all packages OK"
    fi

    return 0
}

detect_cloudways_environment() {
    hostname=$(hostname -f)
    lowercase_hostname=$(echo "$hostname" | tr '[:upper:]' '[:lower:]')

    if echo "$lowercase_hostname" | grep -Eq "\.cloudwaysapps\.com|cloudwaysstagingapps\.com"; then
        is_cloudways=true
    elif [ -f /usr/local/sbin/apm ]; then
        if /usr/local/sbin/apm info | grep -q "Cloudways"; then
            is_cloudways=true
        else
            is_cloudways=false
        fi
    else
        is_cloudways=false
    fi
}

detect_ostype()
{
    if [ ! -f /etc/os-release ]
    then
        OSTYPE=centos
    else
        source /etc/os-release
        if echo $ID $ID_LIKE | grep debian >/dev/null
        then
            OSTYPE=debian
        else
            OSTYPE=centos
        fi
    fi
}

detect_product_debian() {
    if dpkg -l "$PRODUCT_AV" 2>/dev/null | grep -E -q "^.i\s+${PRODUCT_AV}"; then
        PACKAGES="$IMUNIFY_ANTIVIRUS_PACKAGES"
    fi
    if dpkg -l "$PRODUCT_I360" 2>/dev/null | grep -E -q "^.i\s+${PRODUCT_I360}"; then
        PACKAGES="$IMUNIFY360_FIREWALL_PACKAGES"
    fi
}

detect_product_centos() {
    if rpm -q "${PRODUCT_AV}" >/dev/null; then
        PACKAGES="$IMUNIFY_ANTIVIRUS_PACKAGES \
            $IMUNIFY_ANTIVIRUS_RPM_ONLY_PACKAGES"
    fi
    if rpm -q "${PRODUCT_I360}" >/dev/null; then
        PACKAGES="$IMUNIFY360_FIREWALL_PACKAGES \
                 $IMUNIFY360_FIREWALL_RPM_ONLY_PACKAGES"
    fi
}

update_centos()
{
    yum -y update $PACKAGES --enablerepo="imunify360-rollout-*-bypass"
    if rpm -qi imunify-patchman &>/dev/null; then  # optional
        yum -y update imunify-patchman --setopt=obsoletes=0 --enablerepo="imunify360-rollout-*-bypass" || true
    fi
    check_and_reinstall_broken_packages || true
}

update_debian()
{
    if dpkg-query -s imunify-patchman &>/dev/null; then  # optional
        PACKAGES="$PACKAGES imunify-patchman"
    fi
    PACKAGES="$(dpkg-query -W -f='${binary:Package}\n' $PACKAGES 2>/dev/null || true)"
    # first enable rollout bypass repos
    APT_REPO_PATH="/etc/apt/sources.list.d/imunify-rollout-bypass.list"
    UPDATE_PARAMS=''
    cp "${APT_REPO_PATH}.disabled" $APT_REPO_PATH
    apt-get update
    # unhold packages before upgrade to allow update
    apt-mark unhold $PACKAGES 2>/dev/null || true
    # for silent mode
    # force save old config files. new files will be saved by dpkg
    if [ "$DEBIAN_FRONTEND" == "noninteractive" ]; then
       # see dpkg options https://man7.org/linux/man-pages/man1/dpkg.1.html
       # the same options as in the daily update cronjob
       </dev/null apt-get install --only-upgrade -y --allow-change-held-packages -o "Dpkg::Options::=--force-confdef" -o "Dpkg::Options::=--force-confold" $PACKAGES
    else  # interactive
       apt-get install --only-upgrade -y --allow-change-held-packages $PACKAGES
    fi
    # Integrity check runs BEFORE hold/repo-cleanup so reinstall can use bypass repo
    check_and_reinstall_broken_packages || true
    if [ "$is_cloudways" = "true" ]; then
        echo "Cloudways environment detected, holding packages"
        apt-mark hold $PACKAGES 2>/dev/null || true
    fi
    rm $APT_REPO_PATH
}

main() {
    detect_ostype
    detect_product_${OSTYPE}
    detect_cloudways_environment

    if [ -z "$PACKAGES" ]; then
        echo "Neither ImunifyAV nor Imunify360 found installed"
        exit 1
    fi

    update_${OSTYPE}
}

if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
    main "$@"
fi