#!/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=

switched_rw=
cr_name='cr_root'
cr_dev=
blkdev=
mp=
mounted=
tmpdir=$(mktemp -d -t addimageencryption.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

		Encrypt IMAGE

		OPTIONS:
		  --verbose  verbose
		  --prime    add hook scripts to initrd to encrypt on first boot
		  -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
	exit 1
}

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()
{
	if type -p cryptsetup-reencrypt &> /dev/null; then
		echo "$password" | cryptsetup-reencrypt --new "$@"
	else
		echo "$password" | cryptsetup reencrypt --encrypt "$@"
	fi
}

call_dracut()
{
	local initrd="$(readlink "$mp/boot/initrd")"
	local kv="${initrd#initrd-}"
	log_info "create initrd"
	chroot "$mp" dracut --add-drivers dm_crypt -q -f "/boot/$initrd" "$kv" "$@"
}

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
		mountpoint -q "$mp/$i" && continue
		mount -T "$mp"/etc/fstab --target-prefix="$mp" "/$i"
	done
}

make_rw()
{
	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 -n "${0##*/}" -- "$@")
eval set -- "$getopttmp"

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

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

[ -n "$1" ] || helpandquit

if [ -d "$1" ]; then
	mountpoint -q "$1" || err "$1 is not a valid mountpoint"
	mp="$1"
	mounted="$1"
	blkpart="$(findmnt -nvo SOURCE "$mp")"
	[ -L "/sys/class/block/${blkpart##*/}" ] || err "$blkpart is not a partition"
	blkdev="$(readlink -f "/sys/class/block/${blkpart##*/}")"
	blkdev="${blkdev%/*}"
	blkdev="/dev/${blkdev##*/}"
elif [ -b "$1" ]; then
	blkdev="$1"
	blkpart="${blkdev}3"
else
	case "${1##*/}" in
		SLE-Micro.x86_64-5.*-Default-GM.raw )
			log_info "setting up loop device"
			blkdev="$(losetup --show -fP "$1")"
			log_info "loop device $blkdev"
			;;
		openSUSE-MicroOS.x86_64-*-kvm-and-xen*.qcow2)
			[ -e "/dev/nbd0" ] || modprobe nbd
			blkdev=/dev/nbd0
			qemu-nbd -c "$blkdev" "$1"
			udevadm settle
			;;
		*) err "Unsupported image" ;;
	esac
	blkpart="${blkdev}p3"
fi

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 [ -z "$prime" ]; then
	read_password
else
	mkdir -p "$tmpdir/overlay-w"
	dst="$tmpdir/overlay/95addimageencryption"
	mkdir -p "$dst"
	for i in addimageencryption addimageencryption-initrd module-setup.sh \
			addimageencryption-initrd.service; do
		cp "${0%/*}/$i" "$dst/$i"
	done

	make_rw

	mountstuff

	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 3 "$blkdev"
udevadm settle

echo "Encrypting..."
encrypt \
	--type luks1 \
	--reduce-device-size 32m \
	--progress-frequency=1 \
	--iter-time 2000 \
	"${blkpart}"

log_info "Encryption done"

log_info "grow partition again"
echo ", +" | sfdisk -q -N 3 "$blkdev"

log_info "open encrypted image"
echo "$password" | cryptsetup open "${blkpart}" "$cr_name"
cr_dev="/dev/mapper/$cr_name"
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

eval "$(blkid -c /dev/null -o export "$blkpart"|sed 's/^/loop_/')"
echo "$cr_name" "/dev/disk/by-uuid/$loop_UUID" none x-initrd.attach > "$mp"/etc/crypttab

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"
fi

call_dracut

make_ro

echo "Image encryption completed"
