#!/usr/bin/bash

has_template () {
    local arg
    for arg; do
        case $arg in (qubes-template-*) return 0;; esac
    done

    return 1
}

UPDATEVM="$(qubes-prefs --force-root updatevm)"

if [ -z "$UPDATEVM" ]; then
    echo "UpdateVM not set, exiting"
    exit 1
fi

if [ "$1" = "--help" ]; then
    echo "This tool is used to download packages for dom0. Without package list"
    echo "it checks for updates for installed packages"
    echo ""
    echo "Usage: $0 [options] [<pkg list>]"
    echo "    --clean      clean dnf cache before doing anything"
    echo "    --check-only only check for updates (no install)"
    echo "    --gui        use gpk-update-viewer for update selection"
    echo "    --action=... use specific dnf action, instead of automatic install/update"
    echo "    --force-xen-upgrade  force major Xen upgrade even if some qubes are running"
    echo "    --console    does nothing; ignored for backward compatibility"
    echo "    --show-output  does nothing; ignored for backward compatibility"
    echo "    --preserve-terminal  does nothing; ignored for backward compatibility"
    echo "    --skip-boot-check  does not check if /boot & /boot/efi should be mounted"
    echo "    --switch-audio-server-to=(pulseaudio|pipewire) switch audio daemon to pipewire or pulseaudio"
    echo "                 it will be done after requested action (update by default)"
    echo "    <pkg list>   download (and install if run by root) new packages"
    echo "                 in dom0 instead of updating"
    echo "    --           mark end of options"
    echo
    echo "Options with arguments must be passed as -cfoo or --enablerepo=foo,"
    echo "not -c foo or --enablerepo foo"
    exit
fi

PKGS=()
YUM_OPTS=()
UPDATEVM_OPTS=()
QVMTEMPLATE_OPTS=()
GUI=
CHECK_ONLY=
CLEAN=
TEMPLATE=
TEMPLATE_BACKUP=
FORCE_XEN_UPGRADE=
REBOOT_REQUIRED=
DOWNLOADONLY=
AUDIO_SWITCH=
SKIP_BOOT_CHECK=
# Filter out some dnf options and collect packages list
while [ $# -gt 0 ]; do
    case "$1" in
        --enablerepo=*|\
        --disablerepo=*)
            UPDATEVM_OPTS+=( "$1" )
            QVMTEMPLATE_OPTS+=( "$1" )
            ;;
        --downloadonly)
            YUM_OPTS+=( "${1}" )
            QVMTEMPLATE_OPTS+=( "$1" )
            DOWNLOADONLY=1
            ;;
        --clean)
            CLEAN=1
            UPDATEVM_OPTS+=( "$1" )
            ;;
        --gui)
            GUI=1
            UPDATEVM_OPTS+=( "$1" )
            ;;
        --show-output)
            # ignore
            ;;
        --check-only)
            CHECK_ONLY=1
            UPDATEVM_OPTS+=( "$1" )
            ;;
        --preserve-terminal)
            # ignore
            ;;
        --console)
            # ignore
            ;;
        --force-xen-upgrade)
            FORCE_XEN_UPGRADE=1
            ;;
        --switch-audio-server-to=pipewire|\
        --switch-audio-server-to=pulseaudio)
            AUDIO_SWITCH="${1#--switch-audio-server-to=}"
            ;;
        --switch-audio-server-to*)
            echo "Invalid --switch-audio-server-to usage, use either --switch-audio-server-to=pipewire or --switch-audio-server-to=pulseaudio" >&2
            exit 2
            ;;
        --action=*)
            YUM_ACTION=${1#--action=}
            UPDATEVM_OPTS+=( "$1" )
            ;;
        --skip-boot-check)
            SKIP_BOOT_CHECK=1
            ;;
        --)
            if [[ "$#" -gt 1 ]]; then
                YUM_OPTS+=( "${@:2}" )
                UPDATEVM_OPTS+=( "${@:2}" )
                shift "$(( $# - 1 ))"
                : "${YUM_ACTION=install}"
            fi
            ;;
        -*)
            YUM_OPTS+=( "${1}" )
            UPDATEVM_OPTS+=( "$1" )
            QVMTEMPLATE_OPTS+=( "$1" )
            ;;
        *)
            PKGS+=( "${1}" )
            UPDATEVM_OPTS+=( "$1" )
            : "${YUM_ACTION=install}"
            ;;
    esac
    shift
done

REMOTE_ONLY= may_clobber_template=0
case ${YUM_ACTION=upgrade} in
(reinstall|erase|remove|upgrade|upgrade-to|downgrade) may_clobber_template=1;;
(list|search)
    REMOTE_ONLY=1
    ;;
esac

echo "Using $UPDATEVM as UpdateVM for Dom0" >&2

if [ "$CHECK_ONLY" == "1" ]; then
    REMOTE_ONLY=1
    echo "Checking for dom0 updates..." >&2
else
    case ${YUM_ACTION=upgrade} in
    (upgrade|upgrade-to)
        echo "Downloading updates. This may take a while..." >&2;;
    (install|reinstall)
        echo "Downloading packages. This may take a while..." >&2;;
    (downgrade)
        echo "Downgrading packages. This may take a while..." >&2;;
    (list|search)
        echo "Updating package lists. This may take a while..." >&2;;
    (*)
        echo "Performing ${YUM_ACTION}. This may take a while..." >&2;;
    esac
fi

# Redirect operations on templates to qvm-template command
if has_template "${PKGS[@]}"; then
    if [[ ${#PKGS[@]} -ne 1 ]]; then
        echo "ERROR: Specify only one package to reinstall template"
        exit 2
    fi
    if [[ "${PKGS[0]}" =~ ^qubes-template-(.*)(-[0-9]+)(\.[0-9-]+)+(\.[0-9a-zA-Z_]+)?$ ]]; then
        TEMPLATE=${BASH_REMATCH[1]}
    else
        TEMPLATE=${PKGS[0]#qubes-template-}
    fi
    echo "Redirecting to 'qvm-template $YUM_ACTION ${QVMTEMPLATE_OPTS[*]} $TEMPLATE'"
    exec qvm-template "$YUM_ACTION" "${QVMTEMPLATE_OPTS[@]}" "$TEMPLATE"
fi

if [ -n "$AUDIO_SWITCH" ]; then
    if [ -n "$REMOTE_ONLY" ]; then
        echo "--switch-audio-server-to cannot be used with any remote-only action" >&2
        exit 2
    fi
fi

if [ -n "$SWITCHING_AUDIO_IN_PROGRESS" ]; then
    # avoid recursive audio switch
    AUDIO_SWITCH=
fi

# since all the template handling is via qvm-template now, exclude all the templates
TEMPLATE_EXCLUDE_OPTS=( "--exclude=qubes-template-*" )

YUM_OPTS=( "${TEMPLATE_EXCLUDE_OPTS[@]}" "${YUM_OPTS[@]}" )
UPDATEVM_OPTS=( "${TEMPLATE_EXCLUDE_OPTS[@]}" "${UPDATEVM_OPTS[@]}" )

ID=$(id -ur)
if [ "$ID" != 0 ] && [ -z "$GUI" ] && [ -z "$CHECK_ONLY" ] ; then
    echo "This script should be run as root (when used in console mode), use sudo." >&2
    exit 1
fi

LOCKFILE="/var/run/qubes/qubes-dom0-update.lock"
exec 9>"${LOCKFILE}"
[ "$ID" == 0 ] && chgrp qubes "${LOCKFILE}" && chmod g+w "${LOCKFILE}"
flock -n 9
if [[ ${?} -ne 0 ]] && [[ ! $SWITCHING_AUDIO_IN_PROGRESS ]]; then
    echo "Another instance of qubes-dom0-update is already running!"
    exit 1
fi

if [ "$GUI" == "1" ] && [ ${#PKGS[@]} -ne 0 ]; then
    echo "ERROR: GUI mode can be used only for updates" >&2
    exit 1
fi

if [ "$GUI" == "1" ]; then
    apps="xterm konsole yumex apper gpk-update-viewer"

    if [ -n "$KDE_FULL_SESSION" ]; then
        apps="konsole xterm apper yumex gpk-update-viewer"
    fi

    guiapp=
    for app in $apps; do
        if type "$app" &>/dev/null; then
            guiapp=$app
            case $guiapp in
                apper) guiapp="apper --updates --nofork" ;;
                xterm) guiapp="xterm -e sudo dnf update" ;;
                konsole) guiapp="konsole --hold -e sudo dnf update" ;;
                *) guiapp=$app ;;
            esac
            break;
        fi
    done

    if [ -z "$guiapp" ]; then
        message1="You don't have any supported dnf frontend installed."
        message2="Install (using qubes-dom0-update) one of: $apps"

        if [ "$KDE_FULL_SESSION" ]; then
            kdialog --sorry "$message1<br/>$message2"
        else
            zenity --error --text "$message1\n$message2"
        fi

        exit 1
    fi
fi

# Do not start VM automatically when running from cron (only checking for updates)
if [ "$CHECK_ONLY" == "1" ] && ! qvm-check -q --running "$UPDATEVM" > /dev/null 2>&1; then
    echo "ERROR: UpdateVM not running, not starting it in non-interactive mode" >&2
    exit 1
fi

if [ -n "$CLEAN" ]; then
    rm -f /var/lib/qubes/updates/rpm/*
    rm -f /var/lib/qubes/updates/repodata/*
fi
rm -f /var/lib/qubes/updates/errors

# Synopsis: check_mounted MOUNTPOINT
check_mounted() {
    local CHOICE
    # No reason to check further if mount point is already mounted
    awk -v PART="${1}" '{if ($2 == PART ) { exit 0 }} ENDFILE{exit -1}' < /proc/mounts
    [[ ${?} -ne 0 ]] || return
    # No reason to check further if mount point is not in fstab
    awk -v PART="${1}" '!/^[ \t]*#/{ if ( $2 == PART ) { exit 0}} ENDFILE {exit -1}' < /etc/fstab
    [[ ${?} -ne 0 ]] && return
    # Ask user to manually mount partition if user is using GUI Updater
    if [ ! -t 1 ]; then
        echo "Could not decide about unmounted ${1} partition in non-interactive/GUI mode!"
        echo "Please mount ${1} manually before proceeding with updates or update via CLI."
        exit 1
    fi
    read -p "${1} partition is not mounted! mount it now? (y)es, (n)o, (a)bort operation " CHOICE
    case ${CHOICE} in
        y|Y)
            mount "${1}"
            if [[ ${?} -ne 0 ]]; then
                echo "Mounting of ${1} was unsuccessful! aborting."
                exit 1
            fi
            ;;
        n|N) echo "Warning! Proceeding forward without mounting ${1}";;
        a|A) echo Operation aborted!; exit 1;;
        *) echo Invalid choice. Aborting!; exit 1;;
    esac
}

if [ "$SKIP_BOOT_CHECK" != "1" ] && [ "$CHECK_ONLY" != "1" ] && \
        [ "$REMOTE_ONLY" != "1" ] && [ "$DOWNLOADONLY" != "1" ]; then
    # Check if /boot is mounted on split root systems
    check_mounted "/boot"
    # Check if efi partition is mounted on UEFI systems
    [ -d /sys/firmware/efi ] && check_mounted "/boot/efi"
fi

# qvm-run by default auto-starts the VM if not running
readonly dom0_updates_dir=/var/lib/qubes/dom0-updates
rpmdb_path=$(rpm --eval '%{_dbpath}')
rpmdb_path="${rpmdb_path#/}"
qvm-run --nogui -q -u root -- "$UPDATEVM" "mkdir -m 775 -p -- '$dom0_updates_dir'" || exit 1
qvm-run --nogui -q -u root -- "$UPDATEVM" "user=\$(qubesdb-read /default-user) && chown -R -- \"\$user:qubes\" '$dom0_updates_dir'" || exit 1
qvm-run --nogui -q -- "$UPDATEVM" "rm -rf -- '$dom0_updates_dir/etc' '$dom0_updates_dir/$rpmdb_path'" || exit 1
(cd / && exec tar -c -- $rpmdb_path etc/yum/vars etc/yum.repos.d etc/yum.conf etc/dnf/dnf.conf etc/pki/rpm-gpg/RPM-GPG-KEY-* 2>/dev/null) |
   qvm-run --nogui -q --pass-io -- "$UPDATEVM" "LC_MESSAGES=C tar -C '$dom0_updates_dir' -x &&
   sed -Ei 's,^([[:space:]]*gpgkey[[:space:]]*=[[:space:]]*file://)(/etc/pki/rpm-gpg/RPM-GPG-KEY-),\\1$dom0_updates_dir\\2,' '$dom0_updates_dir/etc/yum.repos.d'/*.repo" >/dev/null 2>&1 || {
   status=$?
   echo "Sending repository information to UpdateVM failed: code $status">&2
   exit "$status"
}

CMD="/usr/lib/qubes/qubes-download-dom0-updates.sh --doit --nogui"

# We avoid using bash’s own facilities for this, as they produce $'\n'-style
# strings in certain cases.  These are not portable, whereas the string produced
# by the following is.
for i in "${UPDATEVM_OPTS[@]}"; do CMD+=" '${i//\'/\'\\\'\'}'"; done

QVMRUN_OPTS=(--quiet --filter-escape-chars --nogui --pass-io)
if [[ -t 1 ]] && [[ -t 2 ]]; then
    # Use ‘script’ to emulate a TTY, so that we get status bars and other
    # progress output.  Since stdout and stderr are both terminals, qvm-run
    # will automatically sanitize them, but we explicitly tell it to anyway
    # as a precaution.
    #
    # We MUST NOT use ‘exec script’ here.  That causes ‘script’ to
    # inherit the child processes of the shell.  ‘script’ mishandles
    # this and enters an infinite loop.
    CMD="script --quiet --return --command '${CMD//\'/\'\\\'\'}' /dev/null"
fi

qvm-run "${QVMRUN_OPTS[@]}" -- "$UPDATEVM" "$CMD" < /dev/null

RETCODE=$?
if [[ "$REMOTE_ONLY" = '1' ]] || [ "$RETCODE" -ne 0 ]; then
    if [ "$CHECK_ONLY" = '1' ]; then
        if [ "$RETCODE" -eq 100 ]; then
            echo "There are dom0 updates available" >&2
        elif [ "$RETCODE" -eq 0 ]; then
            echo "No dom0 updates available" >&2
        else
            echo "Failed to check for dom0 updates" >&2
        fi
    fi
    exit $RETCODE
fi
# Wait for download completed
while pidof -x qubes-receive-updates >/dev/null; do sleep 0.5; done

if [ -r /var/lib/qubes/updates/errors ]; then
    echo "*** ERROR while receiving updates:" >&2
    cat /var/lib/qubes/updates/errors >&2
    echo "--> if you want to use packages that were downloaded correctly, use dnf directly now" >&2
    echo "--> otherwise, you can use --clean to remove left-over packages" >&2
    exit 1
fi

# Check for major xen upgrade and warn the user
if [ -f /var/lib/qubes/updates/repodata/repomd.xml ]; then
    xen_in_update=$(ls /var/lib/qubes/updates/rpm/xen-libs-[1-9]* 2>/dev/null \
        |sed -e 's/.*xen-libs-//' \
        |cut -f 1,2 -d . \
        |uniq)
    xen_running=$(xl info xen_version|cut -f 1,2 -d .)
    # cut off -rc if any
    xen_running=${xen_running%-*}
    if [ -n "$xen_in_update" ] && [ "$xen_in_update" != "$xen_running" ]; then
        # check if there are running VMs
        if [ "$(xl list|wc -l)" -gt 2 ]; then
            echo "WARNING: Attempting a major Xen upgrade ($xen_running -> $xen_in_update) while some qubes are running"
            echo "WARNING: You will not be able to interact with them (not even cleanly shutdown) until you restart the system"
            echo "List of running qubes:"
            qvm-ls --running | grep -v dom0
            if [ "$FORCE_XEN_UPGRADE" = 1 ]; then
                echo "Continuing as requested"
            else
                echo -n "Do you want to shutdown all the qubes now? [y/N] "
                read answer
                if [ "$answer" = "y" ] || [ "$answer" = "Y" ]; then
                    qvm-shutdown --all --wait
                else
                    echo "Please shutdown all the qubes, then resume the update with 'sudo dnf upgrade'"
                    exit 1
                fi
            fi
        fi
        if [ -z "$DOWNLOADONLY" ]; then
            REBOOT_REQUIRED=1
        fi
    fi
fi

if [ ${#PKGS[@]} -gt 0 ]; then

    dnf $YUM_ACTION "${YUM_OPTS[@]}" "${PKGS[@]}" ; RETCODE=$?

elif [ -f /var/lib/qubes/updates/repodata/repomd.xml ]; then
    # Above file exists only when at least one package was downloaded
    if [ "$GUI" == "1" ]; then
        # refresh packagekit metadata, GUI utilities use it
        pkcon refresh force
        $guiapp
    else
        dnf check-update ||
        if [ $? -eq 100 ]; then # Run dnf with options
            dnf $YUM_ACTION "${YUM_OPTS[@]}"; RETCODE=$?
        fi
    fi
    if dnf -q check-update; then
        if ! qvm-features dom0 updates-available '' 2>/dev/null; then
            echo "*** WARNING: cannot set feature 'updates-available'" >&2
        fi
        if ! qvm-features dom0 last-update "$(date +'%Y-%m-%d %T')" 2>/dev/null; then
            echo "*** WARNING: cannot set feature 'last-update'" >&2
        fi
    fi
else
    if ! qvm-features dom0 updates-available '' 2>/dev/null; then
        echo "*** WARNING: cannot set feature 'updates-available'" >&2
    fi
    echo "No updates available" >&2
    if [ "$GUI" == "1" ]; then
        if [ "$KDE_FULL_SESSION" ]; then
            kdialog --msgbox 'No updates available'
        else
            zenity --info --title='Dom0 updates' --text='No updates available'
        fi
    fi
fi

# Switching audio

if [ "$AUDIO_SWITCH" = "pipewire" ]; then
    if ! rpm -q pipewire pipewire-pulseaudio >/dev/null; then
        echo "Switching from Pulseaudio to PipeWire" >&2
        SWITCHING_AUDIO_IN_PROGRESS=yes "$0" "${UPDATEVM_OPTS[@]}" --action=install --allowerasing pipewire pipewire-pulseaudio || exit $?
        echo "Audio daemon switched to PipeWire, you can undo it with qubes-dom0-update --switch-audio-server-to=pulseaudio" >&2
        # do not leave the user with stopped both daemons
        user=$(getent group qubes|cut -f 4 -d :|cut -f 1 -d ,) || exit 1
        uid=$(id -u "$user") || exit 1
        sudo -u "$user" XDG_RUNTIME_DIR="/run/user/$uid" systemctl --user start pipewire-pulse
    else
        echo "PipeWire already installed" >&2
    fi
    echo 1 > /var/lib/qubes/.audio-switch-done
elif [ "$AUDIO_SWITCH" = "pulseaudio" ]; then
    if ! rpm -q pulseaudio >/dev/null; then
        echo "Switching from PipeWire to Pulseaudio" >&2
        SWITCHING_AUDIO_IN_PROGRESS=yes "$0" "${UPDATEVM_OPTS[@]}" --allowerasing --action=swap pipewire pulseaudio || exit $?
        echo "Audio daemon switched to Pulseaudio, you can undo it with qubes-dom0-update --switch-audio-server-to=pipewire" >&2
        # do not leave the user with stopped both daemons
        user=$(getent group qubes|cut -f 4 -d :|cut -f 1 -d ,) || exit 1
        uid=$(id -u "$user") || exit 1
        sudo -u "$user" XDG_RUNTIME_DIR="/run/user/$uid" systemctl --user stop pipewire pipewire-pulse pipewire.socket pipewire-pulse.socket
        sudo -u "$user" XDG_RUNTIME_DIR="/run/user/$uid" systemctl --user start pulseaudio
    else
        echo "Pulseaudio already installed" >&2
    fi
    echo 1 > /var/lib/qubes/.audio-switch-done
fi

if [ "$REBOOT_REQUIRED" = 1 ]; then
    echo "This upgrade requires system restart before doing anything else."
    echo -n "Do you want to restart now? [Y/n] "
    read answer
    if [ "$answer" != "n" ] && [ "$answer" != "N" ]; then
        reboot
    fi
fi
exit "$RETCODE"
