#!/bin/bash

# Copyright 2024 NVIDIA Corporation.  All rights reserved.

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}"
    echo -n "Current system partition has "
    grep DISTRIB_DESCRIPTION /etc/lsb-release 2> /dev/null | sed -e 's/DISTRIB_DESCRIPTION=//'
    echo "Other system partition is ${other_part} on ${other_dev}"
    if [ "$(id -u)" != "0" ]; then
        echo "Additional information is available to the root user."
    else
        mount ${other_dev} /mnt 2> /dev/null && {
            other_description=$(grep DISTRIB_DESCRIPTION /mnt/etc/lsb-release 2> /dev/null | sed -e 's/DISTRIB_DESCRIPTION=//')
            umount /mnt
            if [ "${other_description}" != "" ]; then
                echo "Other system partition has ${other_description}"
            fi
        }
        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
}

set_boot_to() {
    /usr/cumulus/bin/grub-switch "$1"
}

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://,,' )
        cp "$image" "$upgrade_installer"
        if [ $? -ne 0 ]; then
            echo "Unable to copy $image" 
            exit 2
        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/partition
    echo "${other_dev}" > $dir/device
    touch $dir/upgrade-launch
    grep -q CL_INSTALLER_UPGRADE_QUICK "$upgrade_installer"
    if [ $? -eq 0 ]; then
        # Quick upgrade without rebooting by installing onto the other
        # partition / device without rebooting; then the next reboot
        # boots directly into the new image.
        $upgrade_installer upgrade "$other_part,$other_dev"
        if [ $? -ne 0 ]; then
            echo "Execution of $upgrade_installer upgrade \"${other_part},${other_dev}\" failed" 
            exit 2
        fi
    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
}

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

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

if [ -z "$activate" -a -z "$cancel" -a -z "$status" -a -z "$upgrade" ]; then
    usage
    exit 0
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 [ "$status" = "yes" ]; then
    show_status
fi

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

if [ "$activate" = "yes" ]; then
    set_boot_to $(other_boot_partition)
fi

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

if [ "$upgrade" = "yes" ]; then
    do_upgrade "$upgrade_image"
fi
