#!/bin/bash
#
# SPDX-FileCopyrightText: Copyright (c) 2023 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.

set -e

EFIVAR_DIR="/sys/firmware/efi/efivars"
PREFIX="L4TOverlayFsMode"
UUID="360a04b9-3fe6-42c8-9fad-d9bf459c4770"
VARNAME="${PREFIX}-${UUID}"
SCRIPT_NAME=$(basename "${0}")
ENABLE_VALUE="0700000001000000"
ENABLE_VALUE_HEX_1="\x07\x00\x00\x00"
ENABLE_VALUE_HEX_2="\x01\x00\x00\x00"
OVERLAYFS_IS_ENABLED=0

usage() {
	echo "Use: ${SCRIPT_NAME} [--enable|-b] [--disable|-d] [--status|-s] [--help|-h]"
cat <<EOF
    This script configures overlayfs for Jetson
    Options are:
    --enable|-e
                   Enable overlayfs
    --disable|-d
                   Disable overlayfs
    --status|-s
                   Query the current overlayfs status
    --help|-h
                   show this help
EOF
}

check_kernel() {
	if [ "$(zcat /proc/config.gz | grep -c 'CONFIG_EFIVAR_FS=y')" -ne 1 ]; then
		echo "CONFIG_EFIVAR_FS must be enabled in the kernel"
		exit 1
	fi

	if [ "$(zcat /proc/config.gz | grep -c 'CONFIG_OVERLAY_FS=y')" -ne 1 ]; then
		echo "CONFIG_OVERLAY_FS must be enabled in the kernel"
		exit 1
	fi
}

check_root() {
	if [ "$(id -u)" -ne 0 ]; then
		echo "This script requires root privilege"
		exit 1
	fi
}

check_efivar() {
	if [[ $(mountpoint "${EFIVAR_DIR}") != *"is a mountpoint"* ]]; then
		echo "${EFIVAR_DIR} not mounted"
		exit 1
	fi
}

remove_var() {
	if [ -f "${EFIVAR_DIR}/${VARNAME}" ]; then
		chattr -i "${EFIVAR_DIR}/${VARNAME}"
		rm "${EFIVAR_DIR}/${VARNAME}"
	fi
}

check_state() {
	# if no varialbe is defined it means overlays was not enabled
	if [ ! -f "${EFIVAR_DIR}/${VARNAME}" ]; then
		OVERLAYFS_IS_ENABLED=0
	else
		val=$(xxd -ps "${EFIVAR_DIR}/${VARNAME}")
		if [ "${val}" = "${ENABLE_VALUE}" ]; then
			OVERLAYFS_IS_ENABLED=1
		else
			OVERLAYFS_IS_ENABLED=0
		fi
	fi
}

status() {
	check_state
	if [ "${OVERLAYFS_IS_ENABLED}" -eq 1 ]; then
		echo "Overlayfs is enabled"
	else
		echo "Overlayfs is disabled"
	fi
}

enable() {
	check_kernel
	check_state
	if [ "${OVERLAYFS_IS_ENABLED}" -eq 1 ]; then
		echo "Overlayfs has already enabled"
		exit 0
	fi
	remove_var

	printf "${ENABLE_VALUE_HEX_1}" > "/tmp/${VARNAME}"
	printf "${ENABLE_VALUE_HEX_2}" >> "/tmp/${VARNAME}"
	dd if="/tmp/${VARNAME}" of="${EFIVAR_DIR}/${VARNAME}" > /dev/null 2>&1
	chattr +i "${EFIVAR_DIR}/${VARNAME}"
	rm "/tmp/${VARNAME}"

	echo "Overlayfs is enabled. Please reboot"
}

disable() {
	if [ ! -f "${EFIVAR_DIR}/${VARNAME}" ]; then
		echo "Overlayfs is not enabled"
	else
		remove_var
		echo "Overlayfs is disabled. Please reboot."
	fi
	exit 0
}

check_root
check_efivar

TGETOPT=$(getopt -n "${SCRIPT_NAME}" --longoptions enable,disable,help,status -o edhs -- "$@")

eval set -- "${TGETOPT}"

while [ $# -gt 0 ]; do
	case "$1" in
	-e|--enable) enable; exit 0 ;;
	-d|--disable) disable; exit 0 ;;
	-s|--status) status; exit 0 ;;
	-h|--help) usage; exit 0 ;;
	--) shift; break ;;
	-*) echo "Unknown option: $@" >&2 ; usage; exit 1 ;;
	esac
	shift
done
