#!/bin/bash

# Copyright 2024-2025 NVIDIA Corporation.  All rights reserved.

is_safe_to_run_upgrade() {
    # Check if any mount is active in /var/installer/sys
    if grep -q "/var/installer/sys" /proc/mounts; then
        echo ""
        echo "/var/installer/sys mounts are present, is $0 -u being run?"
        echo "If not, either reboot or do the following before upgrading:"
        echo "    sudo mount --make-rslave /var/installer/sys"
        echo "    sudo umount -R /var/installer/sys"
        echo "    sudo umount tmpfs-installer"
        echo ""
        return 1
    else
        return 0
    fi
}

is_safe_to_run_activate() {
    # Check if we are running in cold mode
    if ! csmgrctl -s | grep "Current System Mode:" | grep -q "cold"; then
        echo "Error: System is not in cold boot mode (see csmgrctl -s)."
        echo "reboot in cold mode to upgrade"
        return 1
    else
        return 0
    fi
}

has_two_system_partitions() {
    blkid | grep -q 'LABEL="CL-SYSTEM-2"' && echo "yes"
}

current_boot_partition() {
    local rootfs=$(df / | tail -1 | awk '{ print $1 }')
    blkid $rootfs | sed -e 's/^.*CL-SYSTEM-//' -e 's/[^12].*$//'
}

other_boot_partition() {
    local cbp=$( current_boot_partition )
    if [ "$cbp" = '1' ]; then
        echo '2'
    else
        echo '1'
    fi
}

storage_device() {
    lsblk -p | grep disk | awk '{ print $1 }'
}

uuid_of_partition() {
    local label="$1"
    blkid | grep "LABEL=\"$label\"" | sed -e 's/^.* UUID="//' -e 's/".*$//'
}

device_of_partition() {
    local label="$1"
    blkid | grep "LABEL=\"$label\"" | sed -e 's/:.*$//'
}

show_status() {
    current_part=$(current_boot_partition)
    current_dev=$(device_of_partition CL-SYSTEM-$current_part)
    other_part=$(other_boot_partition)
    other_dev=$(device_of_partition CL-SYSTEM-$other_part)
    echo "Current system partition is ${current_part} on ${current_dev}"
    IMAGE_FILE=/etc/image-release
    if [ -f "$IMAGE_FILE" ]; then
        echo -n "Current system partition has "
        IMAGE_RELEASE=$(grep "^IMAGE_RELEASE=" "$IMAGE_FILE" | cut -d "=" -f 2-)
        IMAGE_DESCRIPTION=$(grep "^IMAGE_DESCRIPTION=" "$IMAGE_FILE" | cut -d "=" -f 2-)
        IMAGE_BUILD_SERIAL=$(grep "^IMAGE_BUILD_SERIAL_ID=" "$IMAGE_FILE" | cut -d "=" -f 2-)
        echo "$IMAGE_RELEASE $IMAGE_DESCRIPTION $IMAGE_BUILD_SERIAL"
    else
        other_description=$(grep DISTRIB_DESCRIPTION /etc/lsb-release 2> /dev/null | sed -e 's/DISTRIB_DESCRIPTION=//')
        if [ "${other_description}" != "" ]; then
            echo "Other system partition has ${other_description}"
        fi
    fi

    if [ "$(id -u)" != "0" ]; then
        echo "Additional information is available to the root user."
    else
        if [ -b "${other_dev}" ]; then
            echo "Other system partition is ${other_part} on ${other_dev}"
            mount ${other_dev} /mnt 2> /dev/null && {
                if [ -f /mnt/"$IMAGE_FILE" ]; then
                    IMAGE_RELEASE=$(grep "^IMAGE_RELEASE=" /mnt/"$IMAGE_FILE" | cut -d "=" -f 2-)
                    IMAGE_DESCRIPTION=$(grep "^IMAGE_DESCRIPTION=" /mnt/"$IMAGE_FILE" | cut -d "=" -f 2-)
                    IMAGE_BUILD_SERIAL=$(grep "^IMAGE_BUILD_SERIAL_ID=" /mnt/"$IMAGE_FILE" | cut -d "=" -f 2-)
                    if [ "${IMAGE_RELEASE}" != "" ]; then
                        echo -n "Other system partition has "
                        echo "$IMAGE_RELEASE $IMAGE_DESCRIPTION $IMAGE_BUILD_SERIAL"
                    fi
                else
                    other_description=$(grep DISTRIB_DESCRIPTION /mnt/etc/lsb-release 2> /dev/null | sed -e 's/DISTRIB_DESCRIPTION=//')
                    if [ "${other_description}" != "" ]; then
                        echo "Other system partition has ${other_description}"
                    fi
                fi
                umount /mnt
            }
            next_boot_uuid=$( grep UUID= /boot/grub/grub.cfg | head -1 | sed -e 's/^.*root=UUID=//' -e 's/ .*$//' )
            if [ "$next_boot_uuid" = "$(uuid_of_partition CL-SYSTEM-1)" ]; then
                echo "Next boot to partition 1."
            elif [ "$next_boot_uuid" = "$(uuid_of_partition CL-SYSTEM-2)" ]; then
                echo "Next boot to partition 2."
            else
                echo "Next boot to UUID=$next_boot_uuid"
            fi
        fi
    fi
}

set_boot_to() {
    local partition="$1"
    local rollback="$2"
    local dir="/var/lib/cumulus/image-upgrade"
    mkdir -p "$dir"
    if [ $? -ne 0 ]; then
       echo "Unable to make directory $dir"
       exit 2
    fi

    local other_dev=$(device_of_partition CL-SYSTEM-$partition)
    echo "${other_dev}" > $dir/boot_device

    /usr/cumulus/bin/grub-switch "$partition"

    # Lets enable cumulus-upgrade-on-shutdown service for triggering 
    # backup from running to new target.

    systemctl daemon-reload
    systemctl enable cumulus-upgrade-on-shutdown

    # cumulus-upgrade-on-shutdown uses this file to determine whether
    # to config save/restore (upgrade) or not (rollback).
    if [ "$rollback" != "" ]; then
        echo "rollback" > /var/lib/cumulus/image-upgrade/upgrade-or-rollback
    else
        echo "upgrade" > /var/lib/cumulus/image-upgrade/upgrade-or-rollback
    fi
}

do_upgrade() {
    local image="$1"
    local dir="/var/lib/cumulus/image-upgrade"
    mkdir -p "$dir"
    if [ $? -ne 0 ]; then
       echo "Unable to make directory $dir" 
       exit 2
    fi
    local upgrade_installer=$(mktemp "$dir/upgrade.bin.XXXXXXXX")

    echo "$image" | grep -q -E '^(http(s|)|ftp)://' 
    if [ $? -eq 0 ]; then
        wget -O "$upgrade_installer" "$image"
        if [ $? -ne 0 ]; then
            echo "Unable to fetch $image" 
            exit 2
        fi
    else
        image=$( echo "$image" | sed -e 's,^file://,,' )
        echo "$image" | grep -q '^/var/'
        if [ $? -eq 0 ]; then
            ln -f -s "$image" "$upgrade_installer"
            if [ $? -ne 0 ]; then
                echo "Unable to link $image to $upgrade_installer" 
                exit 2
            fi
        else
            cp "$image" "$upgrade_installer"
            if [ $? -ne 0 ]; then
                echo "Unable to copy $image to $upgrade_installer" 
                exit 2
            fi
        fi
    fi
    check-secure-boot "$upgrade_installer"
    grep -q CL_INSTALLER "$upgrade_installer"
    if [ $? -ne 0 ]; then
        echo "$image does not appear to be Cumulus Linux image" 
        rm -f $"{upgrade_installer}"
        exit 2
    fi
    grep -q CL_INSTALLER_UPGRADE "$upgrade_installer"
    if [ $? -ne 0 ]; then
        echo "$image does not support image upgrade"
        rm -f $"{upgrade_installer}"
        exit 2
    fi
    local other_part=$(other_boot_partition)
    local other_dev=$(device_of_partition CL-SYSTEM-$other_part)
    echo "${other_part}" > $dir/upg_partition
    echo "${other_dev}" > $dir/upg_device
    grep -q CL_INSTALLER_UPGRADE_QUICK "$upgrade_installer"
    if [ $? -eq 0 ]; then
        # First check empty space in /var .  Require at least 2 times
        # the size of the downloaded image or 1.6Gbytes.
        MIN_SPACE_REQUIRED=$(( $(stat -c%s "$upgrade_installer") * 2 ))
        MIN_SPACE_REQUIRED=${MIN_SPACE_REQUIRED:=1500000000}
        if [ $MIN_SPACE_REQUIRED -lt 1600000000 ]; then
            MIN_SPACE_REQUIRED=1600000000
        fi

        # Check available space in /var (given in 1k-blocks, so
        # multiply by 1024 to get it in bytes)
        avail_space=$(( $(df /var | awk 'NR==2 {print $4}' ) * 1024 ))
        avail_space=${avail_space:=0}

        # Check if available space is less than minimum required space
        if [ $avail_space -lt $MIN_SPACE_REQUIRED ]; then
             echo "Error: Available space ${avail_space} in /var is less than ${MIN_SPACE_REQUIRED})."
             echo "Clean up unnecessary files in /var"
             rm $upgrade_installer
             exit 1
        fi

        # Quick upgrade without rebooting by installing onto the other
        # partition / device without rebooting; then the next reboot
        # boots directly into the new image.
        # Changing execute permission for installation
        chmod +x $upgrade_installer
        $upgrade_installer upgrade "$other_part,$other_dev"
        if [ $? -ne 0 ]; then
            echo "Execution of $upgrade_installer upgrade \"${other_part},${other_dev}\" failed"
            rm $upgrade_installer
            exit 2
        fi
        echo "Execution of $upgrade_installer successful"

        rm $upgrade_installer
    else
        # Upgrade that triggers script in
        # initramfs-tools/scripts/local-bottom/cumulus-installer
        # to set up installation of the image through ONIE with the
        # files indicating upgrade to the specific partition / device.
        mv -f $upgrade_installer $dir/upgrade.bin
    fi

    if grep -q "/var/installer/sys" /proc/mounts; then
        # Remove mounts of /sys within /var/installer chroot so that
        # another image upgrade can be done if the user wants to upgrade
        # to a different version.
        mount --make-rslave /var/installer/sys
        umount -R /var/installer/sys
    fi
    if grep -q "tmpfs-installer" /proc/mounts; then
        # tmpfs-installer is a leftover mount from upgrade
        umount tmpfs-installer
    fi
}

run_grub_install() {
    if [ "$1" = "-u" ]; then
        update_grub=yes
    fi
    if [ -d /sys/firmware/efi ]; then
        # workaround for bug #3861745 in generating ONIE entry
        mount /sys/firmware/efi/efivars
        if [ $? -eq 32 ]; then
            efivars_already_mounted=yes
        fi
        mount LABEL="EFI System" /boot/efi
        if [ $? -eq 32 ]; then
            efi_already_mounted=yes
        fi
        grub_args="--uefi-secure-boot"
    else
        # disk device name
        grub_args=$(blkid | grep CL-SYSTEM-1 | sed -e 's,[0-9].*$,,')
    fi
    if [ "$update_grub" = "yes" ]; then
        update-grub
    fi
    grub-install $grub_args
    if [ -d /sys/firmware/efi -a "$efi_already_mounted" != "yes" ]; then
        umount /boot/efi
    fi
    if [ -d /sys/firmware/efi -a "$efivars_already_mounted" != "yes" ]; then
        umount /sys/firmware/efi/efivars
    fi
}

usage() {
    echo "usage: $0 options"
    echo "       -a        activate other image partition at next boot"
    echo "       -c        cancel -a or -r ; next boot will be to current partition"
    echo "       -h        help message (this message)"
    echo "       -r        roll back to other image partition at next boot"
    echo "       -s        show status of system partitions"
    echo "       -u image  upgrade to the specified image in other partition"
}

while getopts "achrsu:" arg; do
    case $arg in
        a)
            activate=yes
            ;;
        c)
            cancel=yes
            ;;
        h)
            usage
            exit 0
            ;;
        r)
            rollback=yes
            ;;
        s)
            status=yes
            ;;
        u)
            upgrade=yes
            upgrade_image="$OPTARG"
            ;;
        *)
            usage
            exit 1
            ;;
    esac
done

if [ -z "$activate" -a -z "$cancel" -a -z "$rollback" -a \
     -z "$status" -a -z "$upgrade" ]; then
    usage
    exit 0
fi

if [ "$status" = "yes" ]; then
    show_status
fi

if [ "$(has_two_system_partitions)" != "yes" ]; then
    echo "Image upgrade can only be done on switches with enough storage"
    echo "to have two system partitions."
    exit 0
fi

if [ "$activate" = "yes" -o "$cancel" = "yes" -o \
     "$rollback" = "yes" -o "$upgrade" = "yes" ]; then
    if [ "$(id -u)" != "0" ]; then
        echo "Must be root to use -a, -c, -r, or -u"
        exit 1
    fi
fi

if [ "$upgrade" = "yes" ]; then
    if ! is_safe_to_run_upgrade; then
	exit 1
    fi

    if [ "$activate" != "yes" ]; then
        run_grub_install -u
    fi
    do_upgrade "$upgrade_image"
    if [ "$activate" != "yes" ]; then
        run_grub_install -u
    fi
fi

if [ "$activate" = "yes" ]; then
    if ! is_safe_to_run_activate; then
	exit 1
    fi
    set_boot_to $(other_boot_partition)
    run_grub_install
fi

if [ "$rollback" = "yes" ]; then
    if ! is_safe_to_run_activate; then
	exit 1
    fi
    set_boot_to $(other_boot_partition) rollback
    run_grub_install
fi


if [ "$cancel" = "yes" ]; then
    set_boot_to $(current_boot_partition)
    run_grub_install
fi

