#!/usr/bin/bash

find_regex_in_args() {
    local regex="${1}"
    shift 1

    for arg in "${@}"; do
        if echo "${arg}" | grep -q -e "${regex}"; then
            return 0
        fi
    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 [--clean] [--check-only] [--gui] [<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, conflicts with --console"
    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    do not open an xterm in the UpdateVM; use this if"
    echo "                 there are GUI problems"
    echo "    --show-output  with --console, show the output of the UpdateVM"
    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=
SHOW_OUTPUT=0
CONSOLE=0
CLEAN=
TEMPLATE=
TEMPLATE_BACKUP=
FORCE_XEN_UPGRADE=
REBOOT_REQUIRED=
# Filter out some dnf options and collect packages list
while [ $# -gt 0 ]; do
    case "$1" in
        --enablerepo=*|\
        --disablerepo=*)
            UPDATEVM_OPTS+=( "$1" )
            QVMTEMPLATE_OPTS+=( "$1" )
            ;;
        --clean)
            CLEAN=1
            UPDATEVM_OPTS+=( "$1" )
            ;;
        --gui)
            GUI=1
            UPDATEVM_OPTS+=( "$1" )
            ;;
        --show-output)
            SHOW_OUTPUT=1
            ;;
        --check-only)
            CHECK_ONLY=1
            ;;
        --console)
            CONSOLE=1
            ;;
        --force-xen-upgrade)
            FORCE_XEN_UPGRADE=1
            ;;
        --action=*)
            YUM_ACTION=${1#--action=}
            UPDATEVM_OPTS+=( "$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

if [[ "$GUI" = 1 ]]; then
    if [[ "$CONSOLE" = 1 ]] || [[ "$SHOW_OUTPUT" = 1 ]]; then
        echo 'Cannot use --gui with --console or --show-output'>&2
        exit 1
    fi
elif [[ -z "${DISPLAY-}" ]]; then
    CONSOLE=1
fi

REMOTE_ONLY= may_clobber_template=0 exit_cmd='&& exit'
case ${YUM_ACTION=upgrade} in
(reinstall|erase|remove|upgrade|upgrade-to|downgrade) may_clobber_template=1;;
(list|search)
    REMOTE_ONLY=1 exit_cmd=
    if [[ "$CONSOLE" = 1 ]]; then SHOW_OUTPUT=1; fi
    ;;
esac

# Redirect operations on templates to qvm-template command
if find_regex_in_args '^qubes-template-' "${PKGS[@]}"; then
    if [[ ${#PKGS[@]} -ne 1 ]]; then
        echo "ERROR: Specify only one package to reinstall template"
        exit 2
    fi
    ONEPKG="$(echo "${PKGS[0]}" | sed -r 's/-[0-9]+(\.[0-9-]+)+(\.noarch)*$//')" # Remove version suffix
    TEMPLATE=${ONEPKG#qubes-template-} # Remove prefix
    echo "Redirecting to 'qvm-template ${YUM_ACTION=install} ${QVMTEMPLATE_OPTS[*]} $TEMPLATE'"
    exec qvm-template "${YUM_ACTION=install}" "${QVMTEMPLATE_OPTS[@]}" "$TEMPLATE"
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

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

echo "Using $UPDATEVM as UpdateVM to download updates for Dom0; this may take some time..." >&2

# qvm-run by default auto-starts the VM if not running
readonly dom0_updates_dir=/var/lib/qubes/dom0-updates
qvm-run --nogui -q -u root -- "$UPDATEVM" "mkdir -m 775 -p -- '$dom0_updates_dir'" || exit 1
qvm-run --nogui -q -u root -- "$UPDATEVM" "chown -R -- user:user '$dom0_updates_dir'" || exit 1
qvm-run --nogui -q -- "$UPDATEVM" "rm -rf -- '$dom0_updates_dir/etc' '$dom0_updates_dir/var/lib/rpm'" || exit 1
(cd / && exec tar -c -- var/lib/rpm 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)
# If we are running in a terminal, run the command in its own terminal in the UpdateVM.
if [[ "$CONSOLE" = 1 ]]; then
    QVMRUN_OPTS+=(--nogui)
    if [[ "$SHOW_OUTPUT" = 1 ]]; then
        QVMRUN_OPTS+=(--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//\'/\'\\\'\'}'"
        fi
    fi
elif [[ "$GUI" = 1 ]] || [[ -t 1 ]] || [[ -t 2 ]]; then
    # In GUI mode, or in a terminal without --console, spawn an xterm to display
    # the output of the remote command.
    CMD="xterm -e /bin/sh -c '"'if "$@" '"$exit_cmd"'; then
    status=0 prompt=succeeded
else
    status=$? prompt="failed with code $?"
fi
read -p "Fetching updates $prompt; press Enter to exit" q
exit "$status"'"' sh $CMD"
else
    QVMRUN_OPTS=(--nogui)
fi

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

RETCODE=$?
if [[ "$REMOTE_ONLY" = '1' ]] || [ "$RETCODE" -ne 0 ]; then
    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
    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 .)
    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
        REBOOT_REQUIRED=1
    fi
fi

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

    dnf "${YUM_OPTS[@]}" $YUM_ACTION "${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_OPTS[@]}" $YUM_ACTION; 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
    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
        zenity --info --title='Dom0 updates' --text='No updates available'
    fi
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"
