#!/bin/bash

CRYPTSETUP="../cryptsetup --type luks2"
REENC=../cryptsetup-reencrypt

DEV_NAME=reenc9768
DEV_NAME2=reenc1273
IMG=reenc-data
ORIG_IMG=reenc-data-orig
KEY1=key1
PWD1="93R4P4pIqAH8"
PWD2="1cND4319812f"
PWD3="1-9Qu5Ejfnqv"

MNT_DIR=./mnt_luks
START_DIR=$(pwd)

function del_scsi_device()
{
	rmmod scsi_debug 2>/dev/null
	sleep 2
}

function remove_mapping()
{
	[ -b /dev/mapper/$DEV_NAME2 ] && dmsetup remove $DEV_NAME2
	[ -b /dev/mapper/$DEV_NAME ] && dmsetup remove $DEV_NAME
	[ ! -z "$LOOPDEV1" ] && losetup -d $LOOPDEV1 >/dev/null 2>&1
	rm -f $IMG $ORIG_IMG $KEY1 >/dev/null 2>&1
	umount $MNT_DIR > /dev/null 2>&1
	rmdir $MNT_DIR > /dev/null 2>&1
	LOOPDEV1=""
	del_scsi_device
}

function fail()
{
	[ -n "$1" ] && echo "$1"
	echo "FAILED at line $(caller)"
	cd $START_DIR
	remove_mapping
	exit 2
}

function skip()
{
	[ -n "$1" ] && echo "$1"
	exit 77
}

function add_scsi_device() {
	del_scsi_device
        modprobe scsi_debug $@
        if [ $? -ne 0 ] ; then
                echo "This kernel seems to not support proper scsi_debug module, test skipped."
                exit 77
        fi

        sleep 2
        SCSI_DEV="/dev/"$(grep -l -e scsi_debug /sys/block/*/device/model | cut -f4 -d /)
        [ -b $SCSI_DEV ] || fail "Cannot find $SCSI_DEV."
}

function open_crypt()
{
	if [ -n "$1" ] ; then
		echo "$1" | $CRYPTSETUP luksOpen $LOOPDEV1 $DEV_NAME || fail
	else
		$CRYPTSETUP luksOpen -d $KEY1 $LOOPDEV1 $DEV_NAME || fail
	fi
}

function wipe_dev() # $1 dev
{
	dd if=/dev/zero of=$1 bs=256k >/dev/null 2>&1
}

function wipe() # $1 pass
{
	open_crypt $1
	wipe_dev /dev/mapper/$DEV_NAME
	udevadm settle >/dev/null 2>&1
	$CRYPTSETUP luksClose $DEV_NAME || fail
}

function prepare() # $1 dev1_siz
{
	remove_mapping

	dd if=/dev/zero of=$IMG      bs=1k count=$1 >/dev/null 2>&1
	LOOPDEV1=$(losetup -f 2>/dev/null)
	[ -z "$LOOPDEV1" ] && fail "No free loop device"
	losetup $LOOPDEV1 $IMG

	if [ ! -e $KEY1 ]; then
		dd if=/dev/urandom of=$KEY1 count=1 bs=32 >/dev/null 2>&1
	fi
}

function check_hash_dev() # $1 dev, $2 hash
{
	HASH=$(sha256sum $1 | cut -d' ' -f 1)
	[ $HASH != "$2" ] && fail "HASH differs ($HASH)"
}

function check_hash() # $1 pwd, $2 hash
{
	open_crypt $1
	check_hash_dev /dev/mapper/$DEV_NAME $2
	$CRYPTSETUP remove $DEV_NAME || fail
}

function backup_orig()
{
	sync
	losetup -d $LOOPDEV1
	cp $IMG $ORIG_IMG
	losetup $LOOPDEV1 $IMG
}

function rollback()
{
	sync
	losetup -d $LOOPDEV1
	cp $ORIG_IMG $IMG
	losetup $LOOPDEV1 $IMG
}

function check_slot() #space separeted list of active key slots
{
	local _out=$($CRYPTSETUP luksDump $LOOPDEV1 | grep -e ": luks2" | sed -e 's/.*\([0-9]\+\):.*/\1/g')

	local _req
	local _hdr
	local _j

	for _i in $*; do
		_j=$((_i))
		_req="$_req $_j"
	done

	for _i in $_out; do
		_j=$((_i))
		_hdr="$_hdr $_j"
	done

	test "$_req" == "$_hdr"
}

function simple_scsi_reenc()
{
	echo -n "$1"
	echo $PWD1 | $CRYPTSETUP luksFormat -i1 $SCSI_DEV || fail

	echo $PWD1 | $CRYPTSETUP luksOpen $SCSI_DEV $DEV_NAME || fail
	HASH=$(sha256sum /dev/mapper/$DEV_NAME | cut -d' ' -f 1)
	$CRYPTSETUP luksClose $DEV_NAME || fail

	echo $PWD1 | $REENC -q -i 1 $SCSI_DEV || fail

	echo $PWD1 | $CRYPTSETUP luksOpen $SCSI_DEV $DEV_NAME || fail
	check_hash_dev /dev/mapper/$DEV_NAME $HASH
	$CRYPTSETUP luksClose $DEV_NAME || fail
}

function mount_and_test() {
	test -d $MNT_DIR || mkdir -p $MNT_DIR
	mount $@ $MNT_DIR 2>/dev/null || {
		echo -n "failed to mount [SKIP]"
		return 0
	}
	rm $MNT_DIR/* 2>/dev/null
	cd $MNT_DIR
	echo $PWD2 | $START_DIR/$REENC $LOOPDEV1 -q --use-fsync --use-directio --write-log || return 1
	cd $START_DIR
	umount $MNT_DIR
	echo -n [OK]
}

function test_logging_tmpfs() {
	echo -n "[tmpfs]"
	mount_and_test -t tmpfs none -o size=$[25*1024*1024] || return 1
	echo
}

function test_logging() {
	echo -n "$1:"
	for img in $(ls img_fs*img.bz2) ; do
		wipefs -a $SCSI_DEV > /dev/null
		echo -n "[${img%.img.bz2}]"
		bzip2 -d -c $img | dd of=$SCSI_DEV >/dev/null 2>&1
		mount_and_test $SCSI_DEV || return 1
	done
	echo
}

[ $(id -u) != 0 ] && skip "WARNING: You must be root to run this test, test skipped."
[ ! -x "$REENC" ] && skip "Cannot find $REENC, test skipped."
which wipefs >/dev/null ||  skip "Cannot find wipefs, test skipped."

# REENCRYPTION tests

HASH4=2daeb1f36095b44b318410b3f4e8b5d989dcc7bb023d1426c492dab0a3053e74
HASH5=bb9f8df61474d25e71fa00722318cd387396ca1736605e1248821cc0de3d3af8
HASH6=4d9cbaf3aa0935a8c113f139691b3daf9c94c8d6c278aedc8eec66a4b9f6c8ae

echo "[1] Reencryption"
prepare 8192
echo $PWD1 | $CRYPTSETUP -q luksFormat -s 128 -c aes-cbc-plain --pbkdf-force-iterations 1 --align-payload 4096 $LOOPDEV1 || fail
wipe $PWD1
check_hash $PWD1 $HASH5
echo $PWD1 | $REENC $LOOPDEV1 -q
check_hash $PWD1 $HASH5
echo $PWD1 | $REENC $LOOPDEV1 -q -s 256
check_hash $PWD1 $HASH5
echo $PWD1 | $REENC $LOOPDEV1 -q -s 256 -c aes-xts-plain64 -h sha256
check_hash $PWD1 $HASH5
echo $PWD1 | $REENC $LOOPDEV1 -q --use-directio
check_hash $PWD1 $HASH5
$CRYPTSETUP luksDump $LOOPDEV1 > /dev/null || fail

echo "[2] Reencryption with data shift"
echo $PWD1 | $CRYPTSETUP -q luksFormat -c aes-cbc-essiv:sha256 -s 128 --pbkdf-force-iterations 1 --align-payload 2048 $LOOPDEV1 || fail
wipe $PWD1
echo $PWD1 | $REENC $LOOPDEV1 -q -s 256 --reduce-device-size 1024S || fail
check_hash $PWD1 $HASH6
echo $PWD1 | $REENC $LOOPDEV1 -q -i 1 || fail
check_hash $PWD1 $HASH6
$CRYPTSETUP luksDump $LOOPDEV1 > /dev/null || fail

echo "[3] Reencryption with keyfile"
echo $PWD1 | $CRYPTSETUP -q luksFormat -d $KEY1 -c aes-cbc-essiv:sha256 -s 128 --pbkdf-force-iterations 1 --align-payload 4096 $LOOPDEV1 || fail
wipe
check_hash "" $HASH5
echo $PWD1 | $CRYPTSETUP -q luksAddKey -d $KEY1 $LOOPDEV1 || fail
$REENC $LOOPDEV1 -d $KEY1 -i 1 -q 2>/dev/null && fail
$REENC $LOOPDEV1 -d $KEY1 -S 0 -i 1 -q || fail
check_hash "" $HASH5
check_slot 0 || fail "Only keyslot 0 expected to be enabled"
$REENC $LOOPDEV1 -d $KEY1 -i 1 -q || fail
$CRYPTSETUP luksDump $LOOPDEV1 > /dev/null || fail
# FIXME echo $PWD1 | $REENC ...

echo "[4] Encryption of not yet encrypted device"
# well, movin' zeroes :-)
OFFSET=8192 # default LUKS2 header size
SIZE=$(blockdev --getsz $LOOPDEV1)
wipe_dev $LOOPDEV1
dmsetup create $DEV_NAME2 --table "0 $(($SIZE - $OFFSET)) linear $LOOPDEV1 0" || fail
check_hash_dev /dev/mapper/$DEV_NAME2 $HASH5
dmsetup remove $DEV_NAME2 || fail
echo $PWD1 | $REENC --type luks2 $LOOPDEV1 -c aes-cbc-essiv:sha256 -s 128 --new --reduce-device-size "$OFFSET"S -q
check_hash $PWD1 $HASH5
$CRYPTSETUP luksDump $LOOPDEV1 > /dev/null || fail

echo "[5] Reencryption using specific keyslot"
echo $PWD2 | $CRYPTSETUP -q luksFormat --pbkdf-force-iterations 1 $LOOPDEV1 || fail
echo -e "$PWD2\n$PWD1" | $CRYPTSETUP -q luksAddKey --pbkdf-force-iterations 1 -S 1 $LOOPDEV1 || fail
echo -e "$PWD2\n$PWD2" | $CRYPTSETUP -q luksAddKey --pbkdf-force-iterations 1 -S 2 $LOOPDEV1 || fail
echo -e "$PWD2\n$PWD1" | $CRYPTSETUP -q luksAddKey --pbkdf-force-iterations 1 -S 3 $LOOPDEV1 || fail
echo -e "$PWD2\n$PWD2" | $CRYPTSETUP -q luksAddKey --pbkdf-force-iterations 1 -S 4 $LOOPDEV1 || fail
echo -e "$PWD2\n$PWD1" | $CRYPTSETUP -q luksAddKey --pbkdf-force-iterations 1 -S 5 $LOOPDEV1 || fail
echo -e "$PWD2\n$PWD2" | $CRYPTSETUP -q luksAddKey --pbkdf-force-iterations 1 -S 6 $LOOPDEV1 || fail
echo -e "$PWD2\n$PWD3" | $CRYPTSETUP -q luksAddKey --pbkdf-force-iterations 1 -S 7 $LOOPDEV1 || fail
backup_orig
echo $PWD2 | $REENC -i 1 -S 0 -q $LOOPDEV1 || fail
check_slot 0 || fail "Only keyslot 0 expected to be enabled"
wipe $PWD2
rollback
echo $PWD1 | $REENC -i 1 -S 1 -q $LOOPDEV1 || fail
check_slot 1 || fail "Only keyslot 1 expected to be enabled"
wipe $PWD1
rollback
echo $PWD2 | $REENC -i 1 -S 6 -q $LOOPDEV1 || fail
check_slot 6 || fail "Only keyslot 6 expected to be enabled"
wipe $PWD2
rollback
echo $PWD3 | $REENC -i 1 -S 7 -q $LOOPDEV1 || fail
check_slot 7 || fail "Only keyslot 7 expected to be enabled"
wipe $PWD3
rollback

echo "[6] Reencryption using all active keyslots"
echo -e "$PWD2\n$PWD1\n$PWD2\n$PWD1\n$PWD2\n$PWD1\n$PWD2\n$PWD3" | $REENC -q $LOOPDEV1 || fail
check_slot 0 1 2 3 4 5 6 7 || fail "All keyslots expected to be enabled"

echo "[7] Reencryption of block devices with different block size"
add_scsi_device sector_size=512 dev_size_mb=8
simple_scsi_reenc "[512 sector]"
add_scsi_device sector_size=4096 dev_size_mb=8
simple_scsi_reenc "[4096 sector]"
add_scsi_device sector_size=512 physblk_exp=3 dev_size_mb=8
simple_scsi_reenc "[4096/512 sector]"
echo "[OK]"

echo "[8] Header only reencryption (hash and iteration time)"
echo $PWD1 | $CRYPTSETUP -q luksFormat --pbkdf-force-iterations 1 --hash sha1 $LOOPDEV1 || fail
wipe $PWD1
check_hash $PWD1 $HASH5
echo $PWD1 | $REENC $LOOPDEV1 -q --keep-key --hash sha256 --iter-time 1
check_hash $PWD1 $HASH5
echo $PWD1 | $REENC $LOOPDEV1 -q --keep-key --hash sha512
check_hash $PWD1 $HASH5
echo $PWD1 | $REENC $LOOPDEV1 -q --keep-key --iter-time 1
check_hash $PWD1 $HASH5
$CRYPTSETUP luksDump $LOOPDEV1 > /dev/null || fail

echo "[9] Test log I/Os on various underlaying block devices"
prepare 8192
echo $PWD2 | $CRYPTSETUP -q luksFormat --pbkdf-force-iterations 1 $LOOPDEV1 || fail
add_scsi_device sector_size=512 dev_size_mb=25
test_logging "[512 sector]" || fail
add_scsi_device sector_size=4096 dev_size_mb=25
test_logging "[4096 sector]" || fail
add_scsi_device sector_size=512 dev_size_mb=25 physblk_exp=3
test_logging "[4096/512 sector]" || fail
test_logging_tmpfs || fail

echo "[10] Removal of encryption"
prepare 8192
echo $PWD1 | $CRYPTSETUP -q luksFormat --pbkdf-force-iterations 1 $LOOPDEV1 || fail
wipe $PWD1
check_hash $PWD1 $HASH5
echo $PWD1 | $REENC $LOOPDEV1 -q --decrypt
check_hash_dev $LOOPDEV1 $HASH4

remove_mapping
exit 0