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

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

verbose=
password=
keyring=
crypttab_options=
root_mp=

is_generated=1
is_rootfs=1
switched_rw=
cr_name=
cr_dev=
blkdev=
blkpart=
mp=
mounted=
keyid=

tmpdir=$(mktemp -d -t disk-encryption-tool.XXXXXX)

cleanup()
{
	set +e
	if [ -n "$mp" ]; then
		while read -r i; do
			[ -z "$is_rootfs" ] || [ "$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
	fi
	[ -d "$tmpdir" ] && ! mountpoint -q "$tmpdir/mnt" && rm -rf "$tmpdir"
}
trap cleanup EXIT

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

		Encrypt MOUNTPOINT or BLOCKDEV.  If not specified, uses /sysroot
		Default VOLUME_NAME is cr_root

		OPTIONS:
		  --verbose       verbose
		  --key           LUKS2 key (use --verbose to print the key if generated)
		  --keyring       store the generated key in the specified keyring
		  --options	  /etc/crypttab options
				  auto: (default) set x-initrd.{mount,attach} if required
				  none: set empty options
				  <options>: set <options>
		  --root          root filesystem (when different from device to encrypt)
		  -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.
	#
	# (see 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
}

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
	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 "$is_generated" ] || echo '{"type":"enrollment-key","keyslots":["0"]}' | cryptsetup token import "${blkpart}"
	else
		cryptsetup reencrypt --batch-mode --verify-passphrase --force-password --verbose --encrypt "${encrypt_options[@]}" "$@" "${blkpart}" "$cr_name"
	fi
	cr_dev="/dev/mapper/$cr_name"
}

make_rw()
{
	local mp="${1:-$mp}"
	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()
{
	local mp="${1:-$mp}"
	[ -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,key:,keyring:,options:,root: -n "${0##*/}" -- "$@")
eval set -- "$getopttmp"

while true ; do
        case "$1" in
                -h|--help) helpandquit ;;
		-v|--verbose) verbose=$((++verbose)); shift ;;
		--key) password="$2"; is_generated=; shift 2 ;;
		--keyring) keyring="$2"; shift 2 ;;
		--options) crypttab_options="$2"; shift 2 ;;
		--root) root_mp="$2"; is_rootfs=; shift 2 ;;
                --) shift ; break ;;
                *) echo "Internal error!" ; exit 1 ;;
        esac
done

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

{ [ -n "$1" ] && [ -n "$2" ]; } || helpandquit

cr_name="$2"
[ -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
fi
shift 2

declare loop_TYPE is_btrfs is_swap
eval "$(blkid -c /dev/null -o export "$blkpart"|sed 's/^/loop_/')"
[ "$loop_TYPE" != crypto_LUKS ] || { echo "Already encrypted"; exit 0; }
[ "$loop_TYPE" != btrfs ] || is_btrfs=1
[ "$loop_TYPE" != swap ] || is_swap=1

if [ -n "$is_btrfs" ]; then
	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

	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
elif [ -n "$is_swap" ]; then
	# sfdisk returns the size in Kilobytes. We choose a very small
	# size, as in any case we need to do the mkswap later again
	minsize=$(($(sfdisk --show-size "$blkpart")*1024))
	minswap=$((512*1024))
	minsize=$((minsize < minswap ? minsize : minswap))
fi

if [ -n "$is_btrfs" ] || [ -n "$is_swap" ]; then
	# Shrink partition to a minimum so reencryption doesn't write
	# everything
	log_info "resizing partition"
	echo "size=$((minsize/1024+32*1024))KiB" | sfdisk --force --no-reread -q -N "$partno" "$blkdev" &> /dev/null
	udevadm settle
	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" || :
	fi
fi

# If a keyring is set, see if the password is stored there and recover
# it.  It is useful for multidisk setups
if [ -n "$keyring" ]; then
	keyid="$(keyctl id %user:"$keyring" 2> /dev/null)" || :
	if [ -n "$keyid" ]; then
		password="$(keyctl pipe "$keyid")"
	fi
fi

if [ -z "$password" ]; then
	password="$(dd if=/dev/urandom bs=8 count=1 2> /dev/null | base64)"
	[ -z "$keyring" ] || echo -n "$password" | keyctl padd user "$keyring" @u > /dev/null
	[ -z "$verbose" ] || echo -e "Enrollment key: \e[1m$password\e[m"
fi

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

if [ -n "$is_btrfs" ] || [ -n "$is_swap" ]; then
	log_info "grow partition again"
	# TODO: recover the size back
	echo ", +" | sfdisk --no-reread -q -N "$partno" "$blkdev" &> /dev/null
	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" || :
		cryptsetup resize "$cr_name" <<<"$password"
	fi
fi

if [ -n "$is_btrfs" ]; then
	if [ -z "$mounted" ]; then
		mount -o rw "$cr_dev" "$tmpdir/mnt"
		mp="$tmpdir/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"

	root_mp="${root_mp:-$mp}"

	make_rw "$root_mp"
elif [ -n "$is_swap" ]; then
	declare loop_UUID
	eval "$(blkid -c /dev/null -o export "$cr_dev"|sed 's/^/loop_/')"
	if [ -n "$loop_UUID" ]; then
		mkswap --uuid "$loop_UUID" "$cr_dev"
	else
		warn "Can't determine device UUID. Can't recreate swap with same UUID"
		mkswap "$cr_dev"
	fi
fi

declare loop_UUID
eval "$(blkid -c /dev/null -o export "$blkpart"|sed 's/^/loop_/')"
if [ -n "$loop_UUID" ]; then
	opts=
	if [ -z "$crypttab_options" ] || [ "$crypttab_options" = "auto" ]; then
		# cr_root and cr_etc are mounted early, in the
		# initrd. In openSUSE cr_var too because of the
		# SELinux re-labeling.  Technically we do not need
		# x-initrd.attach here, as dracut is now able to
		# recognize this.
		#
		# https://systemd.io/MOUNT_REQUIREMENTS/
		if [ "$cr_name" = "cr_root" ]; then
			opts="x-initrd.attach"
		fi
	elif [ "$crypttab_options" != "none" ]; then
		opts="$crypttab_options"
	fi
	[ -n "$opts" ] && echo "$cr_name UUID=$loop_UUID none $opts" >> "$root_mp"/etc/crypttab
	[ -z "$opts" ] && echo "$cr_name UUID=$loop_UUID none" >> "$root_mp"/etc/crypttab
else
	warn "Can't determine device UUID. Can't generate crypttab"
fi

if [ -n "$is_btrfs" ]; then
	make_ro "$root_mp"
fi

echo "Image encryption completed"
