#!/bin/bash

# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#

# This script updates the initrd image by adding modules specified in
# module lists under /etc/nv-update-initrd/list.d.

set -e

init_var()
{
	initrd="/boot/initrd"
	list_dir="/etc/nv-update-initrd/list.d"
	initrd_bak="${initrd}_bak"
	image_paths=("/boot/Image" "/boot/Image.real-time")
	kernel_major_minor_version_pattern="[0-9]+\.[0-9]+"

	if [ ! -f "${initrd}" ]; then
		echo "Error: ${initrd} not found." >&2
		exit 1
	fi

	if [ ! -d "${list_dir}" ]; then
		echo "Error: ${list_dir} not found." >&2
		exit 1
	fi

	local existing_image_paths=()
	# Pre-check for the existence of the kernel image
	for image_path in "${image_paths[@]}"; do
		if [ -f "${image_path}" ]; then
			existing_image_paths+=("${image_path}")
		fi
	done
	if [ "${#existing_image_paths[@]}" -eq 0 ]; then
		echo "Error: No kernel images found!"
		exit 1
	fi
	image_paths=("${existing_image_paths[@]}")

	# The variable is also used in the hook-functions script
	DESTDIR=$(mktemp -d)
	echo "Create the temporary directory ${DESTDIR} for initrd contents. The files will be added later"
}

cleanup()
{
	echo "Cleaning up the temporary directory for updating the initrd.."
	if [ -f "${temp_list}" ]; then
		rm -f "${temp_list}"
	fi
	if [ -d "${DESTDIR}" ]; then
		rm -rf "${DESTDIR}"
	fi
	if [ -f "${list_file}" ]; then
		rm -f "${list_file}"
	fi
	if [ -f "${initrd_bak}" ]; then
		# remove backup initrd
		rm -f "${initrd_bak}"
	fi
}

prepare_list()
{
	local _kernel_version="${1}"
	local _temp_list="${2}"
	local _kernel_major_minor_version=
	local _override_file="modules.override"

	_kernel_major_minor_version=$(get_kernel_major_minor_version "${_kernel_version}")

	# Add the file list that need to be updated to initrd
	if [ -z "${lists_group}"  ]; then
		# Check for modules.override file. If it exists, use ONLY that file for populating initrd.
		if [ -f "${list_dir}/${_override_file}" ]; then
			echo "Exclusively using override file, ${list_dir}/${_override_file}, to populate initrd."
			cat "${list_dir}/${_override_file}" >> "${_temp_list}"
		else
			files="$(find "${list_dir}" -type f,l)"
			for list in ${files}; do
				# Skip files that contain modules and a version number, but don't match the current kernel version
				file_name=$(basename "${list}")
				if [[ "${file_name}" =~ modules.*k${kernel_major_minor_version_pattern} && ! "${file_name}Z" =~ k${_kernel_major_minor_version}Z ]]; then
					continue
				fi

				echo "Including ${list}"
				cat "${list}" >> "${_temp_list}"
			done
		fi
	else
		echo "Using the specified list files"
		old_IFS="${IFS}"
		IFS=','
		for file_name in ${lists_group}; do
			list="${list_dir}/${file_name}"
			if [ ! -f "${list}" ]; then
				echo "Error: ${list} does not exist." >&2
				exit 1
			fi

			echo "Including ${list}"
			cat "${list}" >> "${_temp_list}"
		done
		IFS="${old_IFS}"
	fi

	# Check if list is empty, here we assume that at least one file needs to be updated into initrd
	if [[ ! -s "${_temp_list}" ]]; then
		echo "Error: ${_temp_list} is empty. At least one kernel module or file should to be added to initrd." >&2
		exit 1
	fi
}

modify_initrd()
{
	local _initrd_path="${1}"
	local _kernel_version=
	temp_list=

	pushd "${DESTDIR}" > /dev/null
	gunzip -c "${_initrd_path}" | cpio -i --quiet

	# Remove all files under lib/modules in the initrd and repopulate it later
	# to prevent it from expanding every time the kernel is updated.
	rm -rf lib/modules/*
	# Always explicitly remove /etc/.disable_initrd_bash in the initrd for development.
	# Having a root bash shell launched due to failures results in a security
	# vulnerability. Production systems are expected to package this file.
	rm -f etc/.disable_initrd_bash
	# Add the modules of each image into initrd
	for image_path in "${image_paths[@]}"; do
		# The file list that need to be updated to initrd
		temp_list=$(mktemp)
		_kernel_version=$(get_kernel_version "${image_path}")
		prepare_list "${_kernel_version}" "${temp_list}"

		echo "nv-update-initrd: Updating ${initrd} from ${list_dir} for kernel image ${image_path} and version ${_kernel_version}"
		copy_files_initrd "${DESTDIR}" "${temp_list}" "${_kernel_version}"
		rm -f "${temp_list}"
	done
	# Highlight that production systems should package /etc/.disable_initrd_bash.
	if [ ! -f "etc/.disable_initrd_bash" ]; then
		echo "WARNING: File /etc/.disable_initrd_bash is not present in initrd." \
			"This is appropriate for development but must be packed for production!"
	fi

	# Include config files in the modprobe.d directories for modprobe
	echo "Updating modprobe.d configuration directories for modprobe.."
	MODPROBE_DIRS=("/etc/modprobe.d" "/lib/modprobe.d")
	for modprobe_dir in "${MODPROBE_DIRS[@]}"; do
		# Clean old config files before updating
		rm -rf "${DESTDIR:?}/${modprobe_dir}"

		mkdir -p "${DESTDIR}/${modprobe_dir}"
		for file in "${modprobe_dir}"/*.conf; do
			if test -e "$file"; then
				copy_file config "$file"
			fi
		done
	done

	find . | cpio -H newc -o --quiet | gzip -9 -n > "${_initrd_path}"
	popd > /dev/null
}

copy_files_initrd()
{
	local _initrd_dir="${1}"
	local _list="${2}"
	local _kernel_version="${3}"
	list_file="$(mktemp)"

	# filter out comments and perform variable substitution
	grep -E "^[^#]" "${_list}" | sed \
		-e "s|<KERNEL_VERSION>|${_kernel_version}|g" >"${list_file}"

	local _src=
	local _dst=
	local _dst_dir=
	local _has_error=0
	# Copy all the binary
	while read -r path
	do
		_src="$(echo "${path}" | cut -d ':' -f 1)"
		_dst="$(echo "${path}" | cut -d ':' -f 2)"
		_dst_dir="${_dst%/*}"
		# shellcheck disable=SC2086
		if [ -e "${_src}" ] || find ${_src} > /dev/null 2>&1; then
			echo "Add ${_src} to ${_initrd_dir}/${_dst}"
			mkdir -p "${_initrd_dir}/${_dst_dir}"
			# shellcheck disable=SC2086
			cp -f ${_src} "${_initrd_dir}/${_dst}"
		else
			echo "Error: ${_src} not found." >&2
			_has_error=1
		fi

	done < "${list_file}"

	if [ "${_has_error}" -eq 1 ]; then
		exit 1
	fi
}

get_kernel_version()
{
	local _image_path=${1}
	local _kernel_version=

	# Check if image path is provided
	if [ -z "${_image_path}" ]; then
		echo "Error: No image path provided." >&2
		exit 1
	fi

	# Check if the kernel image exists, which is used to find the kernel version
	if [ -e "${_image_path}" ]; then
		_kernel_version="$(strings "${_image_path}" | sed -nr 's/Linux version ([0-9a-zA-Z.+-]+).*/\1/p' | head -1)"
	fi

	if [ -z "${_kernel_version}" ]; then
		echo "Error: Failed to retrieve kernel version from ${_image_path}." >&2
		exit 1
	fi

	echo "${_kernel_version}"
}

get_kernel_major_minor_version()
{
	# Linux kernel version format: Major.Minor.Patchlevel
	# For example, l4t kernel version:
	# 6.1.9-debug-tegra
	local _kernel_version="${1}"
	local _major_minor_version=

	# Extract the major and minor version numbers
	_major_minor_version=$(echo "${_kernel_version}" | awk -F'[-.]' '{print $1"."$2}')

	# Check if $_major_minor_version is empty or doesn't match the expected format
	if [[ -z "${_major_minor_version}" || ! "${_major_minor_version}" =~ ^${kernel_major_minor_version_pattern}$ ]]; then
		echo "Error: Unable to extract major and minor versions from ${_kernel_version}." >&2
		exit 1
	fi

	echo "${_major_minor_version}"
}

function show_usage
{
	echo "Usage: ${SCRIPT_NAME} [--help|-h] [--list-files|-f]"
cat <<EOF

This script injects the files listed in each specified list file into the initrd image.

	Options are:
		--list-files|-f <list files>
			The group of list files to inject into the initrd image
		--help|-h
			show this help

	Examples:
		1. Use default list files to update the initrd
			${SCRIPT_NAME}
		2. Use the specified list files
			${SCRIPT_NAME} --list-files "modules_common,modules_k6.8"
EOF
}

SCRIPT_NAME=$(basename "$(readlink -f "$0")")

TGETOPT=$(getopt -n "${SCRIPT_NAME}" --longoptions list-files:,help -o f:,h -- "$@")

eval set -- "$TGETOPT"

while [ $# -gt 0 ]; do
	case "$1" in
	-f|--list-files) lists_group="$2"; shift 2 ;;
	-h|--help) show_usage; exit 0 ;;
	--) shift; break ;;
	*) echo "Unknown option: $1" >&2 ; show_usage; exit 0 ;;
	esac
	shift
done

init_var

. /usr/share/nv-update-initrd/hook-functions

# backup initrd
cp "${initrd}" "${initrd_bak}"

# modify backup
trap cleanup EXIT
modify_initrd "${initrd_bak}"

# update initrd
cp "${initrd_bak}" "${initrd}"
