#!/bin/bash

# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: MIT
#
# 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.

#
# gen_luks.sh: Generate LUKS encrypted partition
#
# Usage: gen_luks.sh [--erase <YES>] [--fs-format <ext4>] [--reboot <YES/No>] <dev> <luks_name>
#
# ex. gen_luks.sh /dev/nvme0n1p1 crypt_data
#

erase_data=""
fs_format=""
reboot_sys=""

usage()
{
echo "
Please pass correct parameters.
Usage: sudo gen_luks.sh [--erase <YES>] [--fs-format <ext4>] [--reboot <YES/No>] <dev> <luks_name>
ex. sudo gen_luks.sh /dev/nvme0n1p1 crypt_data"

echo "
options:
--erase <YES>---------- erase the partition before creating encrypted disk.
--fs-format <ext4>----- file system type. Set to ext4 or any other type. If not	ext4, you need to format the partition after device reboot.
--reboot <YES/No>------ reboot the device immediately."

exit 1;
}

# Process options using getopts
while getopts ":-:" OPTION; do
	case $OPTION in
	-) case ${OPTARG} in
		erase)
			erase_data="${!OPTIND}";
			OPTIND=$((OPTIND + 1));
			;;
		fs-format)
			fs_format="${!OPTIND}";
			OPTIND=$((OPTIND + 1));
			;;
		reboot)
			reboot_sys="${!OPTIND}";
			OPTIND=$((OPTIND + 1));
			;;
		*)
			usage;
		;;
		esac;;
	*)
		usage;
		;;
	esac
done

# Shift the processed options and their arguments
shift $((OPTIND - 1))

if [ $# -ne 2 ]; then
	usage;
fi

if cryptsetup isLuks "$1"; then
	echo "${1} is already a LUKS encrypted partition."
	exit 1;
fi

# Check whether <dev> exists or not
enc_dev=$(blkid | grep -w "${1}")
if [ -z "${enc_dev}" ]; then
	echo "No partition ${1} exists, please check again."
	exit 1;
fi

# Check whether <luks_name> is already mounted as an encrypted disk
enc_dev=$(blkid | grep -w "/dev/mapper/${2}")
if [ -n "${enc_dev}" ]; then
	echo "${2} is already being used as a name for encrypted partition."
	echo "Please choose another <luks_name> as second parameter."
	exit 1;
fi

# Check whether <luks_name> is already marked as an encrypted disk
if [ -e "/opt/nvidia/cryptluks" ]; then
	if grep -w "$2" "/opt/nvidia/cryptluks"; then
		echo "${2} is already marked as an encrypted partition."
		echo "Please choose another <luks_name> as second parameter."
		exit 1;
	fi;
fi;

if [ "${erase_data}" != "YES" ]; then
	echo "All data on ${1} will be wiped out after luks disk is created. Reply YES to continue:"
	read confirm
	if [[ "${confirm}" != "YES" ]]; then
		echo "Stop the process."
		exit 1;
	fi;
fi;

# Set default partition format
if [ "${fs_format}" != "ext4" ]; then
	echo -e "\n\nDo you want to format encrypted partition ${2} into ext4? Reply YES or No:"
	read confirm
	if [[ "${confirm}" != "YES" ]]; then
		fs_format="other"
		echo "You need to format encrypted partition /dev/mapper/${2} after device reboot."
	else
		fs_format="ext4"
	fi
fi
# Remove unused partition information
if [ -e "/opt/nvidia/cryptluks" ]; then
	grep -vw "$1" /opt/nvidia/cryptluks > tmp_file
	mv tmp_file /opt/nvidia/cryptluks
fi

# Update encrypted partition table
crypt_disk_uuid=$(cat /proc/sys/kernel/random/uuid)
echo "$1 $2 UUID=${crypt_disk_uuid} ${fs_format}" >> /opt/nvidia/cryptluks

if [ "${reboot_sys}" = "YES" ]; then
	reboot;
elif [ "${reboot_sys}" = "No" ]; then
	echo "You MUST reboot the device manually later to complete the process."
	exit 1;
else
	echo -e "\n\nDo you want to reboot the device to create encrypted partition? Reply YES or No:"
	read confirm
	if [[ "${confirm}" != "YES" ]]; then
		echo "You MUST reboot the device manually later to complete the process."
		exit 1;
	else
		reboot;
	fi;
fi