#!/bin/bash

# SPDX-FileCopyrightText: Copyright (c) 2023-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
#
# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
# property and proprietary rights in and to this material, related
# documentation and any modifications thereto. Any use, reproduction,
# disclosure or distribution of this material and related documentation
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited.

# This script copies the input Capsule payload to the esp and sets the
# "OsIndications" UEFI variable to trigger UEFI Capsule update in the
# next boot.

set -e

BOOT_CTRL_CONF="/etc/nv_boot_control.conf"

ESP_MOUNT_DIR="/boot/efi"
ESP_CAPSULE_DIR="EFI/UpdateCapsule/"
SYS_VAR_DIR="/sys/firmware/efi/efivars"
EFI_GLOBAL_GUID="8be4df61-93ca-11d2-aa0d-00e098032b8c"
TMP_FILE="/tmp/spec_var.bin"
OS_INDICATION_VAR="OsIndications"
UPDATE_CONF="/opt/nvidia/l4t-bootloader-config/nv-l4t-bootloader-config.sh"

# nv_bootloader_capsule_updater.sh "${@}"
function usage()
{
	local script_name=""
	script_name="$(basename "${0}")"

	if [ -n "${1}" ]; then
		echo "${1}"
		echo ""
	fi

	echo "This is a helper script to install firmwares(bootloader components) to corresponding QSPI flash partitions."
	echo ""
	echo "Usage:"
	echo "  ${script_name} [-h|--help] -q <input file>"
	echo "  -h|--help      Displays this help prompt."
	echo ""
	echo "  Positional arguments:"
	echo "  -q <input file> Input QSPI image payload file path."
	echo ""
	echo ""
	echo "Example:"
	echo "  ${script_name} -q /opt/ota_package/Tegra_bl.Cap"
	echo ""

	exit 1
}

function parse_options()
{
	if [ -z "$1" ]; then
		usage "Error: Arguments required"
	fi

	while [ -n "${1}" ]; do
		case "${1}" in
			-h | --help)
				usage
				;;
			-q)
				[ -n "${2}" ] || usage "Not enough parameters"
				qspi_inputfile="${2}"
				shift 2
				;;
			*)
				usage "Error: Unknown option: ${1}"
				;;
		esac
	done

	if [[ "${qspi_inputfile}" == "" ]]; then
		usage "Error: Missed Positional arguments."
	fi
}

function check_prerequisites()
{
	if [ ! -r "${BOOT_CTRL_CONF}" ]; then
		echo "Error: Cannot open ${BOOT_CTRL_CONF} for reading. Exiting..."
		exit 1
	fi

	# Check the payload files and the related utilities
	if [ ! -f "${qspi_inputfile}" ]; then
		echo "Error: Cannot find ${qspi_inputfile}."
		exit 1
	fi

	if [ ! -f "${UPDATE_CONF}" ]; then
		echo "Error: ${UPDATE_CONF} does not exist."
		exit 1
	fi

	if ! "${UPDATE_CONF}" -l; then
		echo "Error: Failed to run ${UPDATE_CONF}."
		exit 1
	fi
}

# Copy the capsule payload to ${esp mount directory}/EFI/UpdateCapsule/
copy_capsule_to_esp () {
	local payload_file="${1}"
	local esp_mount_dir="${2}"
	local capsule_dir=""

	echo "Info: Copy capsule payload: ${payload_file}."
	# Check capsule payload, exit if does not exist
	if [ ! -f "${payload_file}" ]; then
		echo "Error: Cannot find ${payload_file}."
		exit 1
	fi

	capsule_dir="${esp_mount_dir}/${ESP_CAPSULE_DIR}"

	# Create UpdateCapsule directory if does not exist
	if [ ! -d "${capsule_dir}" ]; then
		mkdir -p "${capsule_dir}"
	fi

	# Copy capsule payload and L4Tlauncher
	cp "${payload_file}" "${capsule_dir}"

	echo "Info: Copy capsule payload to ${capsule_dir} done."
}

# Set bit 2 of the OsIndications UEFI variable
function set_capsule_uefi_variable () {
	local var_path=""

	# Set the variable directory
	var_path="${SYS_VAR_DIR}/${OS_INDICATION_VAR}-${EFI_GLOBAL_GUID}"

	if [ -f "${TMP_FILE}" ]; then
		rm -f "${TMP_FILE}"
	fi

	# Set bit 2 of the OsIndications UEFI variable
	printf "\x07\x00\x00\x00\x04\x00\x00\x00\x00\x00\x00\x00" > "${TMP_FILE}"
	dd if="${TMP_FILE}" of="${var_path}" > /dev/null 2>&1

	# Clean up the tmp fie
	rm -f "${TMP_FILE}"

	echo "Info: Set capsule UEFI variable ${var_path} done."
}

# Update the bootloader on QSPI flash.
function update_firmware_on_qspi()
{
	local payload_file="${1}"

	# Since the active esp has been mounted to /boot/efi by the
	# nv-l4t-bootloader-config.sh (called by check_prerequisites()),
	# what is needed here is to check whether the /boot/efi is mount or not.
	if ! mountpoint -q "${ESP_MOUNT_DIR}"; then
		echo "ERROR. The esp is not mounted to ${ESP_MOUNT_DIR}."
		exit 1
	fi

	# Copy the capsule payload to esp directory.
	copy_capsule_to_esp "${payload_file}" "${ESP_MOUNT_DIR}"

	# Set the UEFI variable to indicate UEFI capsule update.
	set_capsule_uefi_variable

	echo "Info: Firmware on QSPI will be updated on the next boot."
}


parse_options "${@}"

check_prerequisites

update_firmware_on_qspi "${qspi_inputfile}"

echo "Info: Reboot the target system for updates to take effect."