#!/bin/bash
#
# Proxmox Multi-Cluster Auto-Update (Debian 13 base)
# ---------------------------------------------------------------
# • Updates nodes grouped by cluster, one cluster fully before the next
# • Cluster with active HA is updated FIRST, the other one afterwards
# • Within a cluster the local node (if it belongs to it) is updated first,
#   then every remote node sequentially
# • Automatically sets up SSH keys if needed
# • Prevents concurrent runs via a lock file
# • Installs needrestart if missing
# • Runs apt update / full-upgrade / autoremove / clean
# • Checks reboot requirement in two ways:
#      – needs-restarting -r  (only reported when a reboot is needed)
#      – /var/run/reboot-required  (used for the e-mail subject)
# • Logs everything and sends an e-mail (no automatic reboot)
# ---------------------------------------------------------------

set -euo pipefail                     # abort on errors / undefined vars

# ------------------- Configuration -------------------
HOSTNAME="$(hostname -f)"               # Fully qualified host name
HOSTNAME_SHORT="$(hostname -s)"
LOCKFILE="/var/run/proxmox-auto-update.lock"
LOGFILE="/var/log/proxmox-cluster-auto-update.log"
MAIL_TO="root@localhost"               # Change to your preferred recipient
APTCMD="/usr/bin/apt-get"              # Absolute path to apt-get
NEEDRESTART_PKG="needrestart"
NEEDRESTART_CMD="/usr/sbin/needs-restarting"
REMOTE_SCRIPT="/tmp/proxmox-remote-update.sh"
SSH_OPTS="-o StrictHostKeyChecking=no -o ConnectTimeout=10"
SSH_KEY="/root/.ssh/id_ed25519"
SSH_PASSWORD=""                        # Same root password for all nodes (optional)
                                       # Leave empty to skip automatic SSH setup

# ---- Cluster groups -------------------------------------------------
# Define which nodes belong to which cluster. Each variable holds a
# space-separated list of hostnames/FQDNs. The LOCAL node does not need
# to be listed – it is detected automatically and, if listed, is updated
# first inside its own cluster.
#
# Example:
#   CLUSTER_A_NODES="pve-a1.lab pve-a2.lab pve-a3.lab"
#   CLUSTER_B_NODES="pve-b1.lab pve-b2.lab"
#
# The cluster that currently has ACTIVE HA resources is always updated
# first. If neither (or both) have HA, CLUSTER_A is updated first.
CLUSTER_A_NODES=""
CLUSTER_B_NODES=""

# Human-readable labels for the log / e-mail
CLUSTER_A_LABEL="Cluster-A"
CLUSTER_B_LABEL="Cluster-B"

# If BOTH lists above are empty the script falls back to single-cluster
# auto-detection (pvecm / /etc/pve/nodes) for a single cluster.
MANUAL_NODES=""
# ---------------------------------------------------------------------

# ----------- Acquire lock (prevent double run) ----------
exec 200>"${LOCKFILE}"
flock -n 200 || {
    echo "$(date '+%Y-%m-%d %H:%M:%S') : another instance is already running on ${HOSTNAME}" \
        >> "${LOGFILE}"
    exit 1
}

# ------------------- Log header -------------------
{
    echo "===================================================================="
    echo "Proxmox Multi-Cluster Auto-Update started: $(date '+%Y-%m-%d %H:%M:%S')"
    echo "Coordinating Node: ${HOSTNAME}"
    echo "Lock file: ${LOCKFILE}"
    echo "Log file: ${LOGFILE}"
    echo "--------------------------------------------------------------------"
} >> "${LOGFILE}"

# ------------------- Error handler -------------------
error_handler() {
    {
        echo "ERROR: The cluster update script terminated on ${HOSTNAME}."
        echo "Please inspect the log file above for details."
    } >> "${LOGFILE}"
    mail -s "[PROXMOX CLUSTER][${HOSTNAME}] Auto-Update failed" "${MAIL_TO}" < "${LOGFILE}"
}
trap error_handler ERR

# ------------------- Install sshpass if needed -------------------
install_sshpass() {
    if [[ -n "${SSH_PASSWORD}" ]] && ! command -v sshpass &> /dev/null; then
        {
            echo "[$(date '+%H:%M:%S')] sshpass not found – installing …"
            ${APTCMD} update -qq
            ${APTCMD} install -y -qq sshpass
            echo "[$(date '+%H:%M:%S')] sshpass installed successfully."
        } >> "${LOGFILE}" 2>&1
    fi
}

# ------------------- Setup SSH key -------------------
setup_ssh_key() {
    {
        echo "[$(date '+%H:%M:%S')] Checking SSH key setup …"
        if [[ ! -f "${SSH_KEY}" ]]; then
            echo "[$(date '+%H:%M:%S')] Generating SSH key …"
            ssh-keygen -t ed25519 -f "${SSH_KEY}" -N "" -C "proxmox-auto-update@${HOSTNAME}"
            echo "[$(date '+%H:%M:%S')] SSH key generated: ${SSH_KEY}"
        else
            echo "[$(date '+%H:%M:%S')] SSH key already exists: ${SSH_KEY}"
        fi
    } >> "${LOGFILE}" 2>&1
}

# ------------------- Copy SSH key to remote node -------------------
copy_ssh_key_to_node() {
    local node=$1
    {
        echo "[$(date '+%H:%M:%S')] Checking SSH access to ${node} …"
        if ssh ${SSH_OPTS} -o BatchMode=yes -o PasswordAuthentication=no "root@${node}" "exit" 2>/dev/null; then
            echo "[$(date '+%H:%M:%S')] SSH key authentication already working for ${node}"
            return 0
        fi
        if [[ -n "${SSH_PASSWORD}" ]]; then
            echo "[$(date '+%H:%M:%S')] SSH key not set up for ${node}, attempting to copy …"
            if ! command -v sshpass &> /dev/null; then
                echo "ERROR: sshpass not installed but SSH_PASSWORD is set"
                return 1
            fi
            if sshpass -p "${SSH_PASSWORD}" ssh-copy-id -o StrictHostKeyChecking=no -i "${SSH_KEY}.pub" "root@${node}" 2>&1; then
                echo "[$(date '+%H:%M:%S')] SSH key successfully copied to ${node}"
                if ssh ${SSH_OPTS} -o BatchMode=yes -o PasswordAuthentication=no "root@${node}" "exit" 2>/dev/null; then
                    echo "[$(date '+%H:%M:%S')] SSH key authentication verified for ${node}"
                    return 0
                else
                    echo "WARNING: SSH key copied but authentication still not working for ${node}"
                    return 1
                fi
            else
                echo "ERROR: Failed to copy SSH key to ${node}"
                return 1
            fi
        else
            echo "WARNING: SSH key not set up for ${node} and no SSH_PASSWORD configured"
            echo "Please either:"
            echo "  1. Set SSH_PASSWORD in the script configuration"
            echo "  2. Manually run: ssh-copy-id root@${node}"
            return 1
        fi
    } >> "${LOGFILE}" 2>&1
}

# ------------------- Auto-detect single cluster nodes -------------------
# Only used as a fallback when CLUSTER_A_NODES and CLUSTER_B_NODES are both empty.
get_cluster_nodes() {
    local nodes=""
    {
        echo "[$(date '+%H:%M:%S')] Detecting cluster nodes (fallback mode) …"
        if command -v pvecm &> /dev/null && pvecm status &> /dev/null 2>&1; then
            while read -r line; do
                if echo "$line" | grep -qE "Nodeid|Membership|---"; then
                    continue
                fi
                local nodename
                nodename=$(echo "$line" | awk '{print $NF}' | sed 's/(local)//')
                if [[ -n "$nodename" ]] && \
                   [[ "$nodename" != "${HOSTNAME}" ]] && \
                   [[ "$nodename" != "${HOSTNAME_SHORT}" ]]; then
                    nodes="${nodes} ${nodename}"
                fi
            done < <(pvecm nodes 2>/dev/null | grep -E "^[[:space:]]*[0-9]")
        fi
        if [[ -z "${nodes}" ]] && [[ -d /etc/pve/nodes ]]; then
            local pve_nodes
            pve_nodes=$(ls -1 /etc/pve/nodes 2>/dev/null | grep -v "^${HOSTNAME_SHORT}\$" || true)
            for node in ${pve_nodes}; do
                local node_fqdn
                node_fqdn=$(getent hosts "${node}" 2>/dev/null | awk '{print $2}' | head -1 || echo "${node}")
                nodes="${nodes} ${node_fqdn}"
            done
        fi
        if [[ -z "${nodes}" ]] && [[ -n "${MANUAL_NODES}" ]]; then
            nodes="${MANUAL_NODES}"
        fi
    } >> "${LOGFILE}" 2>&1
    echo "${nodes}" | xargs
}

# ------------------- Does the local node belong to this list? -------------------
local_in_list() {
    local list="$1"
    for n in ${list}; do
        if [[ "$n" == "${HOSTNAME}" ]] || [[ "$n" == "${HOSTNAME_SHORT}" ]]; then
            return 0
        fi
    done
    return 1
}

# ------------------- Check whether a cluster has ACTIVE HA -------------------
# Returns 0 (has active HA) / 1 (no active HA).
# Queries any reachable node in the list (local shortcut first).
cluster_has_ha() {
    local list="$1"

    # Local shortcut: if this node is part of the cluster, ask locally.
    if local_in_list "${list}"; then
        if command -v ha-manager &> /dev/null; then
            if ha-manager status 2>/dev/null | grep -qE "^service "; then
                return 0
            fi
        fi
        return 1
    fi

    # Otherwise ask the first reachable remote node.
    for node in ${list}; do
        if ssh ${SSH_OPTS} -o BatchMode=yes "root@${node}" \
            "command -v ha-manager >/dev/null 2>&1 && ha-manager status 2>/dev/null | grep -qE '^service '" 2>/dev/null; then
            return 0
        fi
        # If we reached a node and it simply has no HA services, stop asking.
        if ssh ${SSH_OPTS} -o BatchMode=yes "root@${node}" "exit" 2>/dev/null; then
            return 1
        fi
    done
    return 1
}

# ------------------- Create remote update script -------------------
create_remote_script() {
    cat > "${REMOTE_SCRIPT}" << 'REMOTE_EOF'
#!/bin/bash
set -euo pipefail

HOSTNAME="$(hostname -f)"
APTCMD="/usr/bin/apt-get"
NEEDRESTART_PKG="needrestart"
NEEDRESTART_CMD="/usr/sbin/needs-restarting"

echo "[$(date '+%H:%M:%S')] ========== Update starting on ${HOSTNAME} =========="

echo "[$(date '+%H:%M:%S')] Checking for ${NEEDRESTART_PKG} …"
if ! dpkg -s "${NEEDRESTART_PKG}" >/dev/null 2>&1; then
    echo "[$(date '+%H:%M:%S')] ${NEEDRESTART_PKG} missing – installing …"
    ${APTCMD} update -qq
    ${APTCMD} install -y -qq "${NEEDRESTART_PKG}"
    echo "[$(date '+%H:%M:%S')] ${NEEDRESTART_PKG} installed successfully."
else
    echo "[$(date '+%H:%M:%S')] ${NEEDRESTART_PKG} already present."
fi

echo "[$(date '+%H:%M:%S')] apt update …"
${APTCMD} update -qq

echo "[$(date '+%H:%M:%S')] apt full-upgrade -y …"
${APTCMD} full-upgrade -y -qq

echo "[$(date '+%H:%M:%S')] apt autoremove -y …"
${APTCMD} autoremove -y -qq

echo "[$(date '+%H:%M:%S')] apt clean …"
${APTCMD} clean -qq

if ${NEEDRESTART_CMD} -r >/dev/null 2>&1; then
    RESTART_NEEDRESTART="⚠️  needs-restarting: reboot required"
else
    RESTART_NEEDRESTART=""
fi

if [ -f /var/run/reboot-required ]; then
    RESTART_DEBIAN="⚠️  /var/run/reboot-required exists → reboot required"
    SUBJECT_REBOOT_STATUS="Reboot required"
else
    RESTART_DEBIAN="✅  /var/run/reboot-required absent → no reboot needed"
    SUBJECT_REBOOT_STATUS="No reboot needed"
fi

if [[ -n "${RESTART_NEEDRESTART}" ]]; then
    REBOOT_SUMMARY="${RESTART_NEEDRESTART}\n${RESTART_DEBIAN}"
else
    REBOOT_SUMMARY="${RESTART_DEBIAN}"
fi

echo "[$(date '+%H:%M:%S')] ========== Update completed on ${HOSTNAME} =========="
echo -e "${REBOOT_SUMMARY}"
echo "REBOOT_STATUS:${SUBJECT_REBOOT_STATUS}"

exit 0
REMOTE_EOF
    chmod +x "${REMOTE_SCRIPT}"
}

# ------------------- Update local node -------------------
update_local_node() {
    local local_reboot_status="Unknown"
    {
        echo ""
        echo "===================================================================="
        echo "[$(date '+%H:%M:%S')] Updating LOCAL node: ${HOSTNAME}"
        echo "===================================================================="

        echo "[$(date '+%H:%M:%S')] Checking for ${NEEDRESTART_PKG} …"
        if ! dpkg -s "${NEEDRESTART_PKG}" >/dev/null 2>&1; then
            echo "[$(date '+%H:%M:%S')] ${NEEDRESTART_PKG} missing – installing …"
            ${APTCMD} update -qq
            ${APTCMD} install -y -qq "${NEEDRESTART_PKG}"
            echo "[$(date '+%H:%M:%S')] ${NEEDRESTART_PKG} installed successfully."
        else
            echo "[$(date '+%H:%M:%S')] ${NEEDRESTART_PKG} already present."
        fi

        echo "[$(date '+%H:%M:%S')] apt update …"
        ${APTCMD} update -qq

        echo "[$(date '+%H:%M:%S')] apt full-upgrade -y …"
        ${APTCMD} full-upgrade -y -qq

        echo "[$(date '+%H:%M:%S')] apt autoremove -y …"
        ${APTCMD} autoremove -y -qq

        echo "[$(date '+%H:%M:%S')] apt clean …"
        ${APTCMD} clean -qq

        if ${NEEDRESTART_CMD} -r >/dev/null 2>&1; then
            LOCAL_RESTART_NEEDRESTART="⚠️  needs-restarting: reboot required"
        else
            LOCAL_RESTART_NEEDRESTART=""
        fi

        if [ -f /var/run/reboot-required ]; then
            LOCAL_RESTART_DEBIAN="⚠️  /var/run/reboot-required exists → reboot required"
            local_reboot_status="Reboot required"
        else
            LOCAL_RESTART_DEBIAN="✅  /var/run/reboot-required absent → no reboot needed"
            local_reboot_status="No reboot needed"
        fi

        if [[ -n "${LOCAL_RESTART_NEEDRESTART}" ]]; then
            LOCAL_REBOOT_SUMMARY="${LOCAL_RESTART_NEEDRESTART}\n${LOCAL_RESTART_DEBIAN}"
        else
            LOCAL_REBOOT_SUMMARY="${LOCAL_RESTART_DEBIAN}"
        fi

        echo "[$(date '+%H:%M:%S')] Local node update completed"
        echo -e "${LOCAL_REBOOT_SUMMARY}"

        # record result for the summary (inside the redirected block so the
        # variable is available here)
        echo "${HOSTNAME}|SUCCESS|${local_reboot_status}" >> "${RESULTS_FILE}"
    } >> "${LOGFILE}" 2>&1
}

# ------------------- Update remote node -------------------
update_remote_node() {
    local node=$1
    local result_file=$2
    local status="SUCCESS"
    local reboot_status="Unknown"

    {
        echo ""
        echo "===================================================================="
        echo "[$(date '+%H:%M:%S')] Updating REMOTE node: ${node}"
        echo "===================================================================="

        if ! copy_ssh_key_to_node "${node}"; then
            echo "ERROR: Cannot establish SSH connection to ${node}"
            status="FAILED - SSH Setup Failed"
            echo "${node}|${status}|${reboot_status}" >> "${result_file}"
            return 1
        fi

        if ! scp ${SSH_OPTS} "${REMOTE_SCRIPT}" "root@${node}:${REMOTE_SCRIPT}" 2>&1; then
            echo "ERROR: Failed to copy update script to ${node}"
            status="FAILED - Script Transfer Failed"
            echo "${node}|${status}|${reboot_status}" >> "${result_file}"
            return 1
        fi

        local output
        if output=$(ssh ${SSH_OPTS} "root@${node}" "bash ${REMOTE_SCRIPT}" 2>&1); then
            echo "${output}"
            reboot_status=$(echo "${output}" | grep "REBOOT_STATUS:" | cut -d: -f2 || echo "Unknown")
        else
            echo "ERROR: Failed to execute update script on ${node}"
            echo "${output}"
            status="FAILED - Script Execution Failed"
        fi

        ssh ${SSH_OPTS} "root@${node}" "rm -f ${REMOTE_SCRIPT}" 2>&1 || true
        echo "${node}|${status}|${reboot_status}" >> "${result_file}"
    } >> "${LOGFILE}" 2>&1
}

# ------------------- Update one whole cluster -------------------
# $1 = label, $2 = space-separated node list
update_cluster() {
    local label="$1"
    local list="$2"

    {
        echo ""
        echo "####################################################################"
        echo "[$(date '+%H:%M:%S')] >>> Starting cluster: ${label}"
        echo "[$(date '+%H:%M:%S')] Nodes: ${list}"
        echo "####################################################################"
    } >> "${LOGFILE}" 2>&1

    CLUSTER_SUMMARY+="── ${label} ──\n"

    # 1) local node first (only if it is part of this cluster)
    if local_in_list "${list}"; then
        update_local_node
    fi

    # 2) every remote node (skip the local one if it happens to be listed)
    for node in ${list}; do
        if [[ "$node" == "${HOSTNAME}" ]] || [[ "$node" == "${HOSTNAME_SHORT}" ]]; then
            continue
        fi
        update_remote_node "${node}" "${RESULTS_FILE}"
    done

    {
        echo ""
        echo "[$(date '+%H:%M:%S')] <<< Finished cluster: ${label}"
    } >> "${LOGFILE}" 2>&1
}

# ------------------- Main execution -------------------
set +e

RESULTS_FILE="/tmp/proxmox-update-results.$$"
: > "${RESULTS_FILE}"
CLUSTER_SUMMARY="Cluster Update Summary:\n\n"

{
    install_sshpass
    setup_ssh_key
    create_remote_script

    # -------- Multi-cluster mode --------
    if [[ -n "${CLUSTER_A_NODES}" ]] || [[ -n "${CLUSTER_B_NODES}" ]]; then

        echo "[$(date '+%H:%M:%S')] Multi-cluster mode enabled."

        # Determine which cluster has active HA → that one goes first.
        FIRST_LABEL="${CLUSTER_A_LABEL}"; FIRST_NODES="${CLUSTER_A_NODES}"
        SECOND_LABEL="${CLUSTER_B_LABEL}"; SECOND_NODES="${CLUSTER_B_NODES}"

        A_HA="no"; B_HA="no"
        [[ -n "${CLUSTER_A_NODES}" ]] && cluster_has_ha "${CLUSTER_A_NODES}" && A_HA="yes"
        [[ -n "${CLUSTER_B_NODES}" ]] && cluster_has_ha "${CLUSTER_B_NODES}" && B_HA="yes"

        echo "[$(date '+%H:%M:%S')] HA status → ${CLUSTER_A_LABEL}: ${A_HA} | ${CLUSTER_B_LABEL}: ${B_HA}"

        if [[ "${B_HA}" == "yes" ]] && [[ "${A_HA}" != "yes" ]]; then
            # B has HA, A does not → B first
            FIRST_LABEL="${CLUSTER_B_LABEL}"; FIRST_NODES="${CLUSTER_B_NODES}"
            SECOND_LABEL="${CLUSTER_A_LABEL}"; SECOND_NODES="${CLUSTER_A_NODES}"
        fi

        echo "[$(date '+%H:%M:%S')] Update order: 1) ${FIRST_LABEL}  2) ${SECOND_LABEL}"

        [[ -n "${FIRST_NODES}"  ]] && update_cluster "${FIRST_LABEL}"  "${FIRST_NODES}"
        [[ -n "${SECOND_NODES}" ]] && update_cluster "${SECOND_LABEL}" "${SECOND_NODES}"

    # -------- Fallback: plain single-cluster behaviour --------
    else
        echo "[$(date '+%H:%M:%S')] No cluster groups configured – single-cluster fallback."
        update_local_node
        CLUSTER_NODES=$(get_cluster_nodes)
        if [[ -z "${CLUSTER_NODES}" ]]; then
            echo "[$(date '+%H:%M:%S')] Only local node was updated."
        else
            echo "[$(date '+%H:%M:%S')] Found cluster nodes: ${CLUSTER_NODES}"
            for node in ${CLUSTER_NODES}; do
                update_remote_node "${node}" "${RESULTS_FILE}"
            done
        fi
    fi

    # Build the summary block from the results file
    while IFS='|' read -r node status reboot; do
        [[ -z "${node}" ]] && continue
        CLUSTER_SUMMARY+="• ${node}: ${status} - ${reboot}\n"
    done < "${RESULTS_FILE}"

    rm -f "${RESULTS_FILE}"
    rm -f "${REMOTE_SCRIPT}"

} >> "${LOGFILE}" 2>&1

set -e

# ------------------- Log footer & e-mail -------------------
{
    echo "--------------------------------------------------------------------"
    echo "Proxmox Multi-Cluster Auto-Update finished: $(date '+%Y-%m-%d %H:%M:%S')"
    echo "Coordinating Node: ${HOSTNAME}"
    echo ""
    echo -e "${CLUSTER_SUMMARY}"
    echo "===================================================================="
    echo ""
} >> "${LOGFILE}"

mail -s "[PROXMOX CLUSTER] Auto-Update completed on all nodes" "${MAIL_TO}" <<EOF
The scheduled Proxmox multi-cluster update finished.

Clusters were updated one after another; the cluster with active HA was
processed first.

$(echo -e "${CLUSTER_SUMMARY}")

For full details, see the log file at ${LOGFILE} on ${HOSTNAME}.
EOF

exit 0