#!/bin/bash
# SPDX-License-Identifier: MIT
# SPDX-FileCopyrightText: Copyright 2023 SUSE LLC
set -e
shopt -s nullglob

unset "${!LC_@}"
LANG="C.utf8"
export LANG

verbose=
prime=
gen_key=

switched_rw=
cr_name='cr_root'
cr_dev=
blkdev=
blkpart=
partno=
mp=
mounted=
tmpdir=$(mktemp -d -t disk-encryption-tool.XXXXXX)
cleanup()
{
	set +e
	if [ -n "$mp" ]; then
		while read -r i; do
			[ "$i" != "$mp" ] || make_ro
			umount "$i"
		done < <(findmnt -o TARGET -Rn --list "$mp" | tac)
	fi
	if [ -n "$mounted" ]; then
		if [ -e "$tmpdir/mounts" ]; then
			# restore previous mounts
			while read -r line; do
				eval "$line"
				mapfile -td, options < <(echo -n "$OPTIONS")
				if [ -n "$cr_dev" ] && [ "$SOURCE" = "$blkpart" ]; then
					SOURCE="$cr_dev"
				fi
				mount "$SOURCE" "$TARGET" -t "$FSTYPE" -o "$OPTIONS"
			done < "$tmpdir/mounts"
		fi
	else
		[ -e "$cr_dev" ] && cryptsetup close "${cr_dev##*/}"
		case "$blkdev" in
			/dev/nbd*) qemu-nbd -d "$blkdev" ;;
			/dev/loop*) losetup -d "$blkdev" ;;
		esac
	fi
	[ -d "$tmpdir" ] && ! mountpoint -q "$tmpdir/mnt" && rm -rf "$tmpdir"
}
trap cleanup EXIT

helpandquit()
{
	cat <<-EOF
		Usage: $0 [OPTIONS] IMAGE|BLOCKDEV

		Encrypt IMAGE or BLOCKDEV

		OPTIONS:
		  --verbose       verbose
		  --prime         add hook scripts to initrd to encrypt on first boot
		  --partno=N      partition number (default 3)
		  --gen-key       generate random recovery key
		  -h              help screen

	EOF
	exit 0
}

log_info()
{
	[ "${verbose:-0}" -gt 0 ] || return 0
	echo "$@"
}

err()
{
	echo "Error: $*" >&2
	exit 1
}

warn()
{
	echo "Warning: $*" >&2
}

isdigits()
{
       local v="${1:?}"
       [ -z "${v//[0-9]*/}" ]
}

settle_umount_events()
{
        # Manual umount confuses systemd sometimes because it's async and the
        # .mount unit might still be active when the "start" is queued, making
        # it a noop, which ultimately leaves /sysroot unmounted
        # (https://github.com/systemd/systemd/issues/20329). To avoid that,
        # wait until systemd processed the umount events. In a chroot (or with
        # SYSTEMD_OFFLINE=1) systemctl always succeeds, so avoid an infinite loop.
	if [ "$mounted" = "/sysroot" ] && ! systemctl --quiet is-active does-not-exist.mount; then
		while systemctl --quiet is-active sysroot.mount; do sleep 0.5; done
	fi
}

read_password()
{
	local password2
	[ -z "$password" ] || return 0
	if ! [ -t 0 ]; then
		read -r -s password
		return "$?"
	fi
        while true; do
                read -r -s -p "Enter encryption passphrase: " password
		echo
		if type -p pwscore &>/dev/null; then
			echo "$password" | pwscore || continue
		fi
                read -r -s -p "Confirm encryption passphrase: " password2
		echo
                if [ "$password" != "$password2" ]; then
                        echo "Entered passwords don't match"
                        continue
                fi
                [ -n "$password" ] || err "No password, no encryption"
                break
        done
}

encrypt()
{
	local encrypt_options=(\
		--reduce-device-size 32m \
		--progress-frequency=1 \
	)
	if [ -e "${ENCRYPTION_CONFIG:-/etc/encrypt_options}" ]; then
		while read -r op; do
			[ "${op//#}" = "$op" ] || continue
			encrypt_options+=("$op")
		done < "${ENCRYPTION_CONFIG:-/etc/encrypt_options}"
	fi
	# sle 15 compat code
	if type -p cryptsetup-reencrypt &> /dev/null; then
		read_password

		echo "$password" | cryptsetup-reencrypt --new --type luks1 "${encrypt_options[@]}" "$@" "${blkpart}"

		log_info "Encryption done"

		log_info "open encrypted image"
		echo "$password" | cryptsetup open "${blkpart}" "$cr_name"
	else
		log_info "encrypt with options ${encrypt_options[*]}"
		if [ -n "$password" ]; then
			# XXX: hopefully we can use the kernel keyring in the future here
			cryptsetup reencrypt --force-password --verbose --encrypt "${encrypt_options[@]}" "$@" "${blkpart}" "$cr_name" <<<"$password"
			[ -z "$gen_key" ] ||  echo '{"type":"systemd-recovery","keyslots":["0"]}' | cryptsetup token import "${blkpart}"
		else
			cryptsetup reencrypt --batch-mode --verify-passphrase --force-password --verbose --encrypt "${encrypt_options[@]}" "$@" "${blkpart}" "$cr_name"
		fi
	fi
	cr_dev="/dev/mapper/$cr_name"
}

call_dracut()
{
	if [ -L "$mp/boot/initrd" ]; then
		local initrd kv
		initrd="$(readlink "$mp/boot/initrd")"
		kv="${initrd#initrd-}"
		log_info "create initrd for $kv"
		chroot "$mp" dracut --add-drivers dm_crypt -q -f "/boot/$initrd" "$kv" "$@"
	elif [ -e "$mp/etc/kernel/entry-token" ]; then
		local token kernel kv
		read -r token < "$mp/etc/kernel/entry-token"
		log_info "token $token"
		for initrd in "$mp/boot/efi/$token"/*/initrd-*; do
			kv="${initrd%/*}"
			kv="${kv##*/}"
			initrd="${initrd#"$mp"}"
			log_info "create $initrd for $kv"
			hostonly_l=no chroot "$mp" dracut -q --reproducible -f "$initrd" "$kv" "$@"
		done
	else
		err "Unsupported boot loader or fs layout"
	fi
}

mountstuff()
{
	mount -t tmpfs -o size=10m tmpfs "$mp/run"
	for i in proc dev sys; do
		mount --bind "/$i" "$mp/$i"
	done

	for i in /.snapshots /boot/efi /boot/writable /var; do
		grep -q "\s$i\s" "$mp/etc/fstab" || continue
		[ -d "$mp/$i" ] || continue
		mountpoint -q "$mp/$i" && continue
		mount -T "$mp"/etc/fstab --target-prefix="$mp" "/$i"
	done
}

make_rw()
{
	local prop
	read -r prop < <(btrfs prop get -t s "$mp" ro)
	[ "$prop" = "ro=true" ] || return 0
	log_info "switch to rw"
	btrfs prop set -t s "$mp" ro false
	switched_rw=1
}

make_ro()
{
	[ -n "$switched_rw" ] || return 0
	unset switched_rw
	log_info "set ro again"
	btrfs prop set -t s "$mp" ro true
}

####### main #######

getopttmp=$(getopt -o hv --long help,verbose,prime,partno:,gen-key -n "${0##*/}" -- "$@")
eval set -- "$getopttmp"

while true ; do
        case "$1" in
                -h|--help) helpandquit ;;
		-v|--verbose) verbose=$((++verbose)); shift ;;
		--prime) prime="1"; shift ;;
		--partno) partno="$2"; shift 2;;
		--gen-key) gen_key="1"; shift;;
                --) shift ; break ;;
                *) echo "Internal error!" ; exit 1 ;;
        esac
done

[ -z "$1" ] && [ -e /etc/initrd-release ] && set -- /sysroot

[ -n "$1" ] || helpandquit

[ -z "$prime" ] && [ -e "/dev/mapper/$cr_name" ] && err "$cr_name exists. Exit."

if [ -d "$1" ] || [ -b "$1" ]; then
	if [ -b "$1" ]; then
		blkpart="$1"
	else
		mountpoint -q "$1" || err "$1 is not a valid mountpoint"
		mp="$1"
		mounted="$1"
		blkpart="$(findmnt -nvo SOURCE "$mp")"
	fi
	[ -L "/sys/class/block/${blkpart##*/}" ] || err "$blkpart is not a partition"
	blkdev="$(readlink -f "/sys/class/block/${blkpart##*/}")"
	blkdev="${blkdev%/*}"
	blkdev="/dev/${blkdev##*/}"
	read -r partno < "/sys/class/block/${blkpart##*/}"/partition
else
	[ -n "$partno" ] || warn "--partno not specified, assuming 3"
	: "${partno:=3}"
	case "${1##*/}" in
		*.raw )
			log_info "setting up loop device"
			blkdev="$(losetup --show -fP "$1")"
			log_info "loop device $blkdev"
			;;
		*.qcow2)
			[ -e "/dev/nbd0" ] || modprobe nbd
			blkdev=/dev/nbd0
			qemu-nbd -c "$blkdev" "$1"
			udevadm settle
			;;
		*) err "Unsupported image" ;;
	esac
	blkpart="${blkdev}p${partno}"
fi
shift

eval "$(blkid -c /dev/null -o export "$blkpart"|sed 's/^/loop_/')"

[ "$loop_TYPE" != crypto_LUKS ] || { echo "Already encrypted"; exit 0; }
[ "$loop_TYPE" = btrfs ] || err "File system is ${loop_TYPE:-unknown} but only btrfs is supported"

if [ -z "$mounted" ]; then
	log_info "mounting fs"
	mkdir -p "$tmpdir/mnt"
	mount -t btrfs -o rw "${blkpart}" "$tmpdir/mnt"
	mp="$tmpdir/mnt"
else
	mountpoint -q "$mp" || err "$mp is not mounted"
	findmnt -o SOURCE,TARGET,FSTYPE,OPTIONS -Rvn --pairs "$mp" > "$tmpdir/mounts"
	mount -o remount,rw "$mp"
fi

if [ -n "$prime" ]; then
	mkdir -p "$tmpdir/overlay-w"
	dst="$tmpdir/overlay/95disk-encryption-tool"
	mkdir -p "$dst"
	for i in disk-encryption-tool disk-encryption-tool-dracut module-setup.sh \
			disk-encryption-tool-dracut.service generate-recovery-key; do
		cp "${0%/*}/$i" "$dst/$i"
	done
	if [ -e "${ENCRYPTION_CONFIG:-/etc/encrypt_options}" ]; then
		cp "${ENCRYPTION_CONFIG:-/etc/encrypt_options}" "$dst/encrypt_options"
		export ENCRYPTION_CONFIG="/usr/lib/dracut/modules.d/95disk-encryption-tool/encrypt_options"
	fi

	make_rw

	mountstuff

	# dirty
	if [ -d "$mp/usr/share/jeos-firstboot/modules" ]; then
		install -D -m 644 "${0%/*}/jeos-firstboot-diskencrypt-override.conf" "$mp/usr/lib/systemd/system/jeos-firstboot.service.d/jeos-firstboot-diskencrypt-override.conf"
		cp "${0%/*}/jeos-firstboot-enroll" "$mp/usr/share/jeos-firstboot/modules/enroll"
	fi

	mount -t overlay overlay \
		-o lowerdir="$mp/usr/lib/dracut/modules.d/,upperdir=$tmpdir/overlay,workdir=$tmpdir/overlay-w" \
		"$mp/usr/lib/dracut/modules.d/"

	call_dracut

	exit 0
fi

read -r minsize bytes _rest < <(btrfs inspect-internal min-dev-size "$mp")
isdigits "$minsize" || err "Failed to read minimum btrfs size"
[ "$bytes" = 'bytes' ] || err "Failed to read minimum btrfs size"

log_info "resizing fs"
btrfs filesystem resize "$minsize" "$mp"

if [ -e "$tmpdir/mounts" ]; then
	# subshell intentional here
	tac "$tmpdir/mounts" | while read -r line; do
		eval "$line"
		umount "$TARGET"
	done
else
	umount "$mp"
fi
unset mp

settle_umount_events

# shrink partition to a minimum so reencryption doesn't write everything
log_info "resizing partition"
echo "size=$((minsize/1024+32*1024))KiB" | sfdisk -q -N "$partno" "$blkdev"
udevadm settle

if [ -n "$gen_key" ]; then
	read -r password < <(generate-recovery-key)
	echo -e "Recovery key: \e[1m$password\e[m"
	if [ -e /etc/initrd-release ]; then
		read -r key_id < <(echo -n "$password" | keyctl padd user cryptenroll @u)
	fi
fi

echo "Encrypting..."
encrypt "$@"

log_info "grow partition again"
echo ", +" | sfdisk --no-reread -q -N "$partno" "$blkdev"
if [ -e /etc/initrd-release ]; then
	# seems to be the only way to tell the kernel about a specific partition change
	partx -u --nr "$partno" "$blkdev" || :
	# now resize the mapping. For some reason cryptsetup wants a passphrase. Hack
	# around this by installing a token that makes it read the key we installed
	# before, then remove the token again o_O
	cryptsetup token add --key-slot 0 --key-description cryptenroll --token-id 9 "$blkpart"
	cryptsetup resize "$cr_name" < /dev/null
	cryptsetup token remove --token-id 9 "$blkpart"
fi

if [ -z "$mounted" ]; then
	mount -o rw "$cr_dev" "/mnt"
	mp="/mnt"
else
	read -r line < "$tmpdir/mounts"
	eval "$line"
	mapfile -td, options < <(echo -n "$OPTIONS")
	for ((i=0;i<${#options};++i)); do [ "${options[i]}" = ro ] && options[i]=rw; done
	OPTIONS="$(IFS=, eval echo '"${options[*]}"')"
	[ "$SOURCE" = "$blkpart" ] && SOURCE="$cr_dev"
	mount "$cr_dev" "$TARGET" -t "$FSTYPE" -o "$OPTIONS"
	mp="$TARGET"
fi

log_info "resizing fs to max again"
btrfs filesystem resize max "$mp"

make_rw

declare loop_UUID
eval "$(blkid -c /dev/null -o export "$blkpart"|sed 's/^/loop_/')"
if [ -n "$loop_UUID" ]; then
	echo "$cr_name /dev/disk/by-uuid/$loop_UUID none x-initrd.attach" > "$mp"/etc/crypttab
else
	warn "Can't determine device UUID. Can't generate crypttab"
fi

mountstuff

if grep -q "LOADER_TYPE.*grub2" "$mp"/etc/sysconfig/bootloader; then
	log_info "Update bootloader"

	echo GRUB_ENABLE_CRYPTODISK=y >> "$mp"/etc/default/grub

	sed -i -e 's/^LOADER_TYPE=.*/LOADER_TYPE="grub2"/' "$mp"/etc/sysconfig/bootloader
	chroot "$mp" update-bootloader --reinit
	sed -i -e 's/^LOADER_TYPE=.*/LOADER_TYPE="grub2-efi"/' "$mp"/etc/sysconfig/bootloader
	chroot "$mp" update-bootloader --reinit
	mv "$mp/boot/grub2/grub.cfg" "$mp/boot/grub2/grub.cfg.bak"
	cat > "$mp/boot/grub2/grub.cfg" <<-'EOF'
	set linux=linux
	set initrd=initrd
	if [ "${grub_cpu}" = "x86_64" -o "${grub_cpu}" = "i386" ]; then
	    if [ "${grub_platform}" = "efi" ]; then
		set linux=linuxefi
		set initrd=initrdefi
	    fi
	fi
	export linux initrd
	EOF
	sed -e 's/linuxefi/$linux/;s/initrdefi/$initrd/' < "$mp/boot/grub2/grub.cfg.bak" >> "$mp/boot/grub2/grub.cfg"
	rm "$mp/boot/grub2/grub.cfg.bak"
elif [ -e "$mp/etc/kernel/entry-token" ]; then
	: # sd-boot
else
	log_info "unsupported boot loader"
	grep LOADER_TYPE "$mp"/etc/sysconfig/bootloader
fi

# A new initrd is created as side effect of the enrolment
# (jeos-firtboot module), as this calls sdbootutil

make_ro

echo "Image encryption completed"
