#!/bin/bash

# SPDX-FileCopyrightText: Copyright (c) 2024 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.

#
# nv_create_usbkey.sh: Create USB key for auto disk unlock
#
# Usage: nv_create_usbkey.sh <usb_dev> <enc_dev>
#
# ex. sudo nv_create_usbkey.sh /dev/sda3 /dev/nvme0n1p3
#

set -e
set -o pipefail

usage()
{
echo "
Please pass correct parameters.
Usage: nv_create_usbkey.sh <usb_dev> <enc_dev>
ex. sudo nv_create_usbkey.sh /dev/sda3 /dev/nvme0n1p3"

exit 1;
}

cleanup() {
	# Clean up mounted device
	if [ -e "${usb_dev}" ]; then
		if findmnt "${usb_dev}" >/dev/null; then
			umount "${usb_dev}"
			rmdir "${usb_mountpath}" >/dev/null
		fi
	fi
}
trap cleanup EXIT

# Process options using getopts
while getopts ":-:" OPTION; do
	case $OPTION in
	*)
		usage;
		;;
	esac
done

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

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

usb_dev="${1}"
enc_dev="${2}"

# If the tools are installed
if [ ! -e "/usr/sbin/cryptsetup" ]; then
	echo "Please install cryptsetup debian package."
	exit 1;
fi

if [ ! -e "/bin/findmnt" ]; then
	echo "Please install util-linux debian package."
	exit 1;
fi

# If the device is LUKS encypted
if ! cryptsetup isLuks "${enc_dev}"; then
	echo "${enc_dev} is NOT a LUKS encrypted partition."
	exit 1;
fi

# If the device already has 2 keys
# One is user password and the other is USB key
key_num=$(cryptsetup luksDump "${enc_dev}" | grep -c ": luks2")
if [ "${key_num}" != 1 ]; then
	echo "USB key already created before.";
	exit 1;
fi

# Mount point USB storage
if findmnt "${usb_dev}" >/dev/null; then
	umount "${usb_dev}"
fi
usb_mountpath="/mnt/tmp"
mkdir -p "${usb_mountpath}"
mount "${usb_dev}" "${usb_mountpath}"

# Create USB key and add this key into the LUKS encrypted disk
uuid=$(blkid -o value -s UUID "${enc_dev}")
key_path="${usb_mountpath}/${uuid}.lek"
dd if=/dev/urandom bs=1 count=256 | tee "${key_path}" > /dev/null

if ! cryptsetup luksAddKey "${enc_dev}" "${key_path}"; then
	echo "Add USB key failed"
	exit 1;
fi

# Create USB unlock script
usb_unlock_script="/bin/usb_unlock"
cat << "END" > ${usb_unlock_script}
#!/bin/sh

set -e

if [ ! -e "/mnt/unlock" ]; then
	mkdir -p "/mnt/unlock"
fi

for usbpartition in /dev/disk/by-id/usb-*; do
	usbdevice=$(readlink -f ${usbpartition})
		if mount ${usbdevice} "/mnt/unlock" 2>/dev/null; then
			if [ -e "/mnt/unlock/${CRYPTTAB_KEY}.lek" ]; then
				cat "/mnt/unlock/${CRYPTTAB_KEY}.lek"
				umount ${usbdevice} 2>/dev/null
				rmdir "/mnt/unlock" 2>/dev/null
				exit 0
			fi
			umount ${usbdevice} 2>/dev/null
		fi
done

rmdir "/mnt/unlock" 2>/dev/null
/lib/cryptsetup/askpass "Insert USB key or input passphrase and press ENTER:"
END

chmod 755 ${usb_unlock_script}

# Set uuid for key path searching into field 3
# Set keyscript path and initramfs into field 4
awk -i inplace -v uuid="${uuid}" -v uuid_str="UUID=${uuid}" \
	-v option="luks,keyscript=${usb_unlock_script},initramfs" \
	'$2 == uuid_str { $3 = uuid; $4 = option } 1' /etc/crypttab

# copy USB unlock script into initramfs
initramfs_copy_unlock_script="/etc/initramfs-tools/hooks/unlock-from-usb"
cat << "END" > ${initramfs_copy_unlock_script}
#!/bin/sh

PREREQ=""

prereqs() {
        echo "$PREREQ"
}

case "$1" in
        prereqs)
                prereqs
                exit 0
        ;;
esac

. "${CONFDIR}/initramfs.conf"
. /usr/share/initramfs-tools/hook-functions

# Copy key search script from system to bin folder under initramfs
copy_exec "/bin/usb_unlock" "/bin"
END

# Update initramfs
update-initramfs -k all -u

# Remove USB unlock script since it was copied into initramfs
rm ${usb_unlock_script}