#!/bin/sh
# Copyright (C) 2022 Fondries.IO
# SPDX-License-Identifier: MIT
#
# Encrypt (reencrypt) root device with LUKS2

fatal() {
	echo "$1"
	exit 1
}

msg() {
	echo "$1"
}

cryptfs_enabled() {
	return 0
}

fd_check() {
	if [ ! -d "/dev/fd" ]; then
		`/bin/ln -s /proc/self/fd /dev/fd`
	fi
}

e2fsck_check() {
	fsckret=0
	if [ -n "`which e2fsck`" ]; then
		fsckout=`e2fsck -p -v ${1}`
		fsckret=$?
		# Avoid empty newline after summary
		echo "e2fsck: ${fsckout}" >/dev/kmsg
		# Return code >= 4 means uncorrected / operational error
		## TODO: force boot into a recovery mode or similar, as there is really not
		## much we can do in case the fs is corrupted in a bad way
		if [ "${fsckret}" -ge "4" ]; then
			echo "e2fsck: WARNING: file system errors left uncorrected: ret ${fsckret}" >/dev/kmsg
		fi
	fi
	return "${fsckret}"
}

cryptfs_gen_passphrase() {
	# Static as at this point we just need a key for encrypting and later enrolling a new keyslot
	mkdir -p /run/cryptsetup
	echo -n "scle" > /run/cryptsetup/passphrase
}

cryptfs_run() {
	# Similar to rootfs, we need to wait for the device to become available
	C=0
	delay=${bootparam_datadelay:-1}
	timeout=${bootparam_datatimeout:-5}
	while true; do
		if [ $(( $C * $delay )) -gt $timeout ]; then
			fatal "root '$bootparam_data' doesn't exist or does not contain a /dev."
		fi

		if [ -n "$bootparam_data" ]; then
			root_dev="$bootparam_data"
			if [ "`echo ${bootparam_data} | cut -c1-5`" = "UUID=" ]; then
				root_uuid=`echo $bootparam_data | cut -c6-`
				root_dev=`readlink -f /dev/disk/by-uuid/$root_uuid`
			elif [ "`echo ${bootparam_data} | cut -c1-9`" = "PARTUUID=" ]; then
				root_partuuid=`echo $bootparam_data | cut -c10-`
				root_dev=`readlink -f /dev/disk/by-partuuid/$root_partuuid`
			elif [ "`echo ${bootparam_data} | cut -c1-10`" = "PARTLABEL=" ]; then
				root_partlabel=`echo $bootparam_data | cut -c11-`
				root_dev=`readlink -f /dev/disk/by-partlabel/$root_partlabel`
			elif [ "`echo ${bootparam_data} | cut -c1-6`" = "LABEL=" ]; then
				root_label=`echo $bootparam_data | cut -c7-`
				root_dev=`readlink -f /dev/disk/by-label/$root_label`
			fi

			[ -e "$root_dev" ] && break
		fi
		debug "Sleeping for $delay second(s) to wait root to settle..."
		sleep $delay
		C=$(( $C + 1 ))
	done

	flags=""
	#root_dev="/dev/mmcblk0p3"
	mounted_dir="/data"
	key_slot=8

	# Identify desired token format (e.g. pkcs11, tpm2, etc) and import required functions
	if [ ! -d /etc/cryptfs ]; then
		fatal "No initramfs cryptfs module found"
	fi
	luks_token=`ls /etc/cryptfs | head -n1`
	if [ -z "${luks_token}" ]; then
		fatal "No valid initramfs cryptfs module found"
	fi

	if [ ! -d "${mounted_dir}" ]; then
		/bin/mkdir -p "${mounted_dir}"
	fi

	fd_check

	. /etc/cryptfs/${luks_token}

	cryptfs_check_${luks_token}

	cryptfs_gen_passphrase

	if ! cryptsetup isLuks ${root_dev}; then
		# Partition not yet encrypted
		msg "${root_dev} not yet encrypted, encrypting with LUKS2"
		e2fsck_check ${root_dev}
		block_size=`dumpe2fs -h ${root_dev} 2>/dev/null | grep "^Block size" | cut -d ':' -f 2 | tr -d ' '`
		block_count=`dumpe2fs -h ${root_dev} 2>/dev/null | grep "^Block count" | cut -d ':' -f 2 | tr -d ' '`
		luks_size=33554432 # 32M
		new_block_count=$(($block_count - $luks_size / $block_size))
		resize2fs -p ${root_dev} ${new_block_count}
		if [ $? -ne 0 ]; then
			fatal "Failed to resize ${root_dev} to allow extra size required for luks support"
		fi

		cat /run/cryptsetup/passphrase | cryptsetup -v luksFormat --type luks2 --key-slot ${key_slot} --disable-locks --reduce-device-size 32m ${root_dev}

		# Align label and UUID if used as boot parameter (not safe, better use the proper device path instead)
		if [ -n "$root_label" ]; then
			cryptsetup config --label ${root_label} ${root_dev}
		fi
		if [ -n "$root_uuid" ]; then
			yes | cryptsetup luksUUID --uuid ${root_uuid} ${root_dev}
		fi
	fi

	luks_name="`basename ${root_dev}`_crypt"

	# Check if online encryption is still in progress
	if cryptsetup luksDump ${root_dev} | grep -q "online-reencrypt"; then
		# Run recovery process
		cat /run/cryptsetup/passphrase | cryptsetup luksOpen ${root_dev} ${luks_name}
		e2fsck_check /dev/mapper/${luks_name}
		cat /run/cryptsetup/passphrase | cryptsetup -v reencrypt --resume-only ${root_dev}
		cryptsetup close ${luks_name}
	fi

	cryptfs_pre_${luks_token}

	if ! cryptsetup luksDump ${root_dev} | grep -q "clevis"; then
		msg "Enrolling LUKS2 keyslot based on ${luks_token} token"
		cryptfs_enroll_${luks_token} ${root_dev} ${key_slot}
	fi

	cryptfs_post_${luks_token} ${root_dev} ${luks_name}

	e2fsck_check /dev/mapper/${luks_name}
	if [ $? -gt "0" ]; then
		/sbin/mkfs.ext4 -q -j /dev/mapper/${luks_name}
	fi

	mount ${flags} /dev/mapper/${luks_name} "$ROOTFS_DIR/data" ||
		(cryptsetup luksClose ${luks_name} && fatal "Failed to mount LUKS ${luks_name}")
}