#!/bin/bash

# Copyright (c) 2020-2021, NVIDIA CORPORATION. All rights reserved.
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
# DEALINGS IN THE SOFTWARE.

# This script:
#
# Read in update_list.conf and store the partition list under “Partitions” into an
# array.
#
# Parse “Slot” setting and convert it into corresponding suffix by retrieving current slot suffix
# through calling nvbootctrl.
#
# Repeatedly call nv_bootloader_payload_updater by passing in one partition name a time:
# Use option --part <part_name>, Form <part_name> by concatenating name from “Partitions” list and
# suffix from the result of above converting “Slot” setting.

set -e
set -o pipefail

DEFAULT_UPDATE_CONF="update_list.conf"

usage()
{
    echo -e "
Usage: $0 [path]
Where,
    path    Indicate the path of update_list.conf file (Optional)

This script reads in an update partition list, get the corresponding suffix from nvbootctrl and call
nv_bootloader_payload_updater to update individual partitions specified in the list.
    "; echo;
    exit 1
}

get_slot_suffix()
{
    local slot="${1}"
    local part_name="${2}"
    local current_slot
    local non_current_slot
    local slot_a_suffix
    local slot_b_suffix
    local current_slot_suffix
    local non_current_slot_suffix
    local ret_code

    # Determine whether ROOTFS A/B and unified A/B is enabled by checking
    # the return code of the command "nvbootctrl is-unified-enabled"
    # EX_OK(0): Rootfs AB is enabled with unified AB.
    # EX_SOFTWARE(70): Rootfs AB is enabled without unified AB.
    # EX_UNAVAILABLE(69): Rootfs AB isn't enabled.
    set +e
    "${_nvbootctrl}" is-unified-enabled >/dev/null 2>&1
    ret_code=$?
    set -e
    if [ "${ret_code}" == 0 ] || [ "${ret_code}" == 69 ]; then
        # If unified A/B is enabled or rootfs A/B is not enabled, find curernt
        # slot using bootloader
        current_slot=$(${_nvbootctrl} get-current-slot)
    elif [ "${ret_code}" == 70 ]; then
        # if rootfs a/b is enabled but unified a/b is not enabled, find current
        # slot using rootfs if part_name is kernel/kernel-dtb
        if [[ "${part_name}" = "kernel" || "${part_name}" = "kernel-dtb" ]]; then
            current_slot=$(${_nvbootctrl} -t rootfs get-current-slot)
        else
            # Otherwise, search normally
            current_slot=$(${_nvbootctrl} get-current-slot)
        fi
    else
        echo "Failed to run \"${_nvbootctrl} is-unified-enabled\". Exiting..."
        return 1
    fi

    non_current_slot=$((1 - current_slot))
    slot_a_suffix=$(${_nvbootctrl} get-suffix "0" | tail -1)
    slot_b_suffix=$(${_nvbootctrl} get-suffix "1" | tail -1)
    current_slot_suffix="$(${_nvbootctrl} get-suffix "${current_slot}" | tail -1)"
    non_current_slot_suffix="$(${_nvbootctrl} get-suffix "${non_current_slot}" | tail -1)"
    case "${slot}" in
        "current") slot_suffixes=("${current_slot_suffix}") ;;
        "non-current") slot_suffixes=("${non_current_slot_suffix}") ;;
        "a") slot_suffixes=("${slot_a_suffix}") ;;
        "b") slot_suffixes=("${slot_b_suffix}") ;;
        "both")
            slot_suffixes=("${slot_a_suffix}" "${slot_b_suffix}")
        ;;
    esac
}

partition_update_request()
{
    local lines_num
    local item

    echo "Update config file is ${update_conf}"
    readarray update_conf_array < "${update_conf}"

    lines_num=${#update_conf_array[@]}
    echo "Number of lines is $lines_num"

    max_index=$((lines_num - 1))
    echo "max_index=${max_index}"

    for i in $(seq 0 ${max_index})
    do
        item=${update_conf_array[$i]}

        # Skip the empty line
        if [ "${#item}" == 1 ]; then
            continue
        fi
        if echo "${item}" | grep -q "^#"; then
            continue
        fi

        part_name=$(echo "${item}" | awk '{print $1}')
        slot=$(echo "${item}" | awk '{print $2}')
        declare -a slot_suffixes
        get_slot_suffix "${slot}" "${part_name}"
        for j in "${slot_suffixes[@]}"
        do
            echo "Slot for ${part_name} is ${j}"
            echo "Running: ${_nv_bootloader_payload_updater} --part ${part_name}${j}"
            ${_nv_bootloader_payload_updater} --part "${part_name}${j}"
        done

    done
}


prerequisite_check()
{

    if ! which nvbootctrl;then
        echo "nvbootctrl is not found"
        exit 1
    fi
    _nvbootctrl=$(which nvbootctrl)

    if ! which nv_bootloader_payload_updater;then
        echo "nv_bootloader_payload_updater is not found"
        exit 1
    fi

    _nv_bootloader_payload_updater=$(which nv_bootloader_payload_updater)

    if [ "${USER}" != "root" ]; then
        echo "${0} requires root privilege";
        exit 1;
    fi
}


opstr="h"
while getopts "${opstr}" OPTION; do
    case $OPTION in
    *)
       usage
       ;;
    esac;
done


if [ $# -eq 1 ]; then
    update_conf=${1};
fi;

if [ -z "${update_conf}" ]; then
    update_conf="${DEFAULT_UPDATE_CONF}"
fi

if ! [ -f "${update_conf}" ]; then
    echo "${update_conf}: No such file"
    if [ $# -eq 0 ]; then
        usage
    fi;
    exit 1
fi

prerequisite_check
partition_update_request

