#!/bin/bash

[ -z "$CRYPTSETUP_PATH" ] && CRYPTSETUP_PATH=".."
CRYPTSETUP=$CRYPTSETUP_PATH/cryptsetup
REENC=$CRYPTSETUP_PATH/cryptsetup-reencrypt
FAST_PBKDF="--pbkdf-force-iterations 1000"

DEV_NAME=reenc9768
DEV_NAME2=reenc1273
IMG=reenc-data
IMG_HDR=$IMG.hdr
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 --retry $DEV_NAME2
	[ -b /dev/mapper/$DEV_NAME ] && dmsetup remove --retry $DEV_NAME
	[ ! -z "$LOOPDEV1" ] && losetup -d $LOOPDEV1 >/dev/null 2>&1
	rm -f $IMG $IMG_HDR $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 backtrace:"
	while caller $frame; do ((frame++)); done
	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 $@ delay=0
        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() # $1 pwd, $2 hdr
{
	if [ -n "$2" ] ; then
		echo "$1" | $CRYPTSETUP luksOpen $LOOPDEV1 $DEV_NAME --header $2 || fail
	elif [ -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, $3 hdr
{
	open_crypt $1 $3
	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 separated list of ENABLED key slots
{
	local _KS0=DISABLED
	local _KS1=$_KS0 _KS2=$_KS0 _KS3=$_KS0 _KS4=$_KS0 _KS5=$_KS0 _KS6=$_KS0 _KS7=$_KS0
	local _tmp

	for _tmp in $*; do
		eval _KS$_tmp=ENABLED
	done

	local _out=$($CRYPTSETUP luksDump $LOOPDEV1 | grep -e "Key Slot" | cut -d ' ' -f 4)

	local _i=0
	for _tmp in $_out; do
		eval local _orig="\${_KS${_i}}"
		if [ "$_tmp" != "$_orig" ]; then
			echo "Keyslot $_i is $_tmp, expected result: $_orig"
			return 1
		fi
		_i=$[_i+1]
	done

	return 0
}

function simple_scsi_reenc()
{
	echo -n "$1"
	echo $PWD1 | $CRYPTSETUP luksFormat --type luks1 $FAST_PBKDF $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 $FAST_PBKDF $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

	if [ "${REENC:0:1}" != "/" ] ; then
		MNT_REENC=$START_DIR/$REENC
	else
		MNT_REENC=$REENC
	fi

	echo $PWD2 | $MNT_REENC $LOOPDEV1 -q --use-fsync --use-directio --write-log $FAST_PBKDF || 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.xz) ; do
		wipefs -a $SCSI_DEV > /dev/null
		echo -n "[${img%.img.xz}]"
		xz -d -c $img | dd of=$SCSI_DEV bs=4k >/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 2>&1 ||  skip "Cannot find wipefs, test skipped."

# REENCRYPTION tests

HASH1=b69dae56a14d1a8314ed40664c4033ea0a550eea2673e04df42a66ac6b9faf2c
HASH2=d85ef2a08aeac2812a648deb875485a6e3848fc3d43ce4aa380937f08199f86b
HASH3=e4e5749032a5163c45125eccf3e8598ba5ed840df442c97e1d5ad4ad84359605
HASH4=2daeb1f36095b44b318410b3f4e8b5d989dcc7bb023d1426c492dab0a3053e74
HASH5=5f70bf18a086007016e948b04aed3b82103a36bea41755b6cddfaf10ace3c6ef

echo "[1] Reencryption"
prepare 8192
echo $PWD1 | $CRYPTSETUP -q luksFormat --type luks1 -s 128 -c aes-cbc-plain $FAST_PBKDF --align-payload 4096 $LOOPDEV1 || fail
wipe $PWD1
check_hash $PWD1 $HASH1
echo $PWD1 | $REENC $LOOPDEV1 -q $FAST_PBKDF
check_hash $PWD1 $HASH1
echo $PWD1 | $REENC $LOOPDEV1 -q -s 256 $FAST_PBKDF
check_hash $PWD1 $HASH1
echo $PWD1 | $REENC $LOOPDEV1 -q -s 256 -c aes-xts-plain64 -h sha256 $FAST_PBKDF
check_hash $PWD1 $HASH1
echo $PWD1 | $REENC $LOOPDEV1 -q --use-directio $FAST_PBKDF
check_hash $PWD1 $HASH1
echo $PWD1 | $REENC $LOOPDEV1 -q --master-key-file /dev/urandom $FAST_PBKDF
check_hash $PWD1 $HASH1
echo $PWD1 | $REENC $LOOPDEV1 -q -s 512 --master-key-file /dev/urandom $FAST_PBKDF
check_hash $PWD1 $HASH1
$CRYPTSETUP --type luks1 luksDump $LOOPDEV1 > /dev/null || fail

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

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

echo "[4] Encryption of not yet encrypted device"
# well, movin' zeroes :-)
OFFSET=2048
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 $HASH3
dmsetup remove --retry $DEV_NAME2 || fail
echo $PWD1 | $REENC $LOOPDEV1 -c aes-cbc-essiv:sha256 -s 128 --new --type luks1 --reduce-device-size "$OFFSET"S -q $FAST_PBKDF || fail
check_hash $PWD1 $HASH3
$CRYPTSETUP --type luks1 luksDump $LOOPDEV1 > /dev/null || fail
# 64MiB + 1 KiB
prepare 65537
OFFSET=131072
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 --retry $DEV_NAME2 || fail
echo $PWD1 | $REENC $LOOPDEV1 -c aes-cbc-essiv:sha256 -s 128 --new --type luks1 --reduce-device-size "$OFFSET"S -q $FAST_PBKDF || fail
check_hash $PWD1 $HASH5
$CRYPTSETUP --type luks1 luksDump $LOOPDEV1 > /dev/null || fail
prepare 8192
OFFSET=4096
echo fake | $REENC $LOOPDEV1 -d $KEY1 --new --type luks1 --reduce-device-size "$OFFSET"S -q $FAST_PBKDF || fail
$CRYPTSETUP open --test-passphrase $LOOPDEV1 -d $KEY1 || fail
wipe_dev $LOOPDEV1

echo "[5] Reencryption using specific keyslot"
echo $PWD2 | $CRYPTSETUP -q luksFormat --type luks1 $FAST_PBKDF $LOOPDEV1 || fail
echo -e "$PWD2\n$PWD1" | $CRYPTSETUP -q luksAddKey $FAST_PBKDF -S 1 $LOOPDEV1 || fail
echo -e "$PWD2\n$PWD2" | $CRYPTSETUP -q luksAddKey $FAST_PBKDF -S 2 $LOOPDEV1 || fail
echo -e "$PWD2\n$PWD1" | $CRYPTSETUP -q luksAddKey $FAST_PBKDF -S 3 $LOOPDEV1 || fail
echo -e "$PWD2\n$PWD2" | $CRYPTSETUP -q luksAddKey $FAST_PBKDF -S 4 $LOOPDEV1 || fail
echo -e "$PWD2\n$PWD1" | $CRYPTSETUP -q luksAddKey $FAST_PBKDF -S 5 $LOOPDEV1 || fail
echo -e "$PWD2\n$PWD2" | $CRYPTSETUP -q luksAddKey $FAST_PBKDF -S 6 $LOOPDEV1 || fail
echo -e "$PWD2\n$PWD3" | $CRYPTSETUP -q luksAddKey $FAST_PBKDF -S 7 $LOOPDEV1 || fail
backup_orig
echo $PWD2 | $REENC $FAST_PBKDF -S 0 -q $LOOPDEV1 || fail
check_slot 0 || fail "Only keyslot 0 expected to be enabled"
wipe $PWD2
rollback
echo $PWD1 | $REENC $FAST_PBKDF -S 1 -q $LOOPDEV1 || fail
check_slot 1 || fail "Only keyslot 1 expected to be enabled"
wipe $PWD1
rollback
echo $PWD2 | $REENC $FAST_PBKDF -S 6 -q $LOOPDEV1 || fail
check_slot 6 || fail "Only keyslot 6 expected to be enabled"
wipe $PWD2
rollback
echo $PWD3 | $REENC $FAST_PBKDF -S 7 -q $LOOPDEV1 || fail
check_slot 7 || fail "Only keyslot 7 expected to be enabled"
wipe $PWD3
rollback
echo $PWD3 | $REENC $FAST_PBKDF -S 8 -q $LOOPDEV1 2>/dev/null && fail
$CRYPTSETUP luksDump $LOOPDEV1 > /dev/null || fail

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 $FAST_PBKDF || 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 --type luks1 --hash sha1 $FAST_PBKDF $LOOPDEV1 || fail
wipe $PWD1
check_hash $PWD1 $HASH1
echo $PWD1 | $REENC $LOOPDEV1 -q --keep-key || fail
check_hash $PWD1 $HASH1
echo $PWD1 | $REENC $LOOPDEV1 -q --keep-key --pbkdf-force-iterations 999 2>/dev/null && fail
check_hash $PWD1 $HASH1
echo $PWD1 | $REENC $LOOPDEV1 -q --keep-key --hash sha256 --pbkdf-force-iterations 1001
check_hash $PWD1 $HASH1
[ "$($CRYPTSETUP luksDump $LOOPDEV1 | grep -A1 -m1 "Key Slot 0" | grep Iterations: | sed -e 's/[[:space:]]\+Iterations:\ \+//g')" -eq 1001 ] || fail
[ "$($CRYPTSETUP luksDump $LOOPDEV1 | grep -m1 "Hash spec:" | cut -f2)" = "sha256" ] || fail
echo $PWD1 | $REENC $LOOPDEV1 -q --keep-key --hash sha512 $FAST_PBKDF
check_hash $PWD1 $HASH1
[ "$($CRYPTSETUP luksDump $LOOPDEV1 | grep -A1 -m1 "Key Slot 0" | grep Iterations: | sed -e 's/[[:space:]]\+Iterations:\ \+//g')" -eq 1000 ] || fail
echo $PWD1 | $REENC $LOOPDEV1 -q --keep-key $FAST_PBKDF
check_hash $PWD1 $HASH1
$CRYPTSETUP --type luks1 luksDump $LOOPDEV1 > /dev/null || fail

echo "[9] Test log I/Os on various underlaying block devices"
prepare 8192
echo $PWD2 | $CRYPTSETUP -q luksFormat --type luks1 $FAST_PBKDF $LOOPDEV1 || fail
add_scsi_device sector_size=512 dev_size_mb=32
test_logging "[512 sector]" || fail
add_scsi_device sector_size=4096 dev_size_mb=32
test_logging "[4096 sector]" || fail
add_scsi_device sector_size=512 dev_size_mb=32 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 --type luks1 $FAST_PBKDF $LOOPDEV1 || fail
wipe $PWD1
check_hash $PWD1 $HASH1
echo $PWD1 | $REENC $LOOPDEV1 -q --decrypt || fail
check_hash_dev $LOOPDEV1 $HASH4

echo $PWD1 | $CRYPTSETUP -q luksFormat --type luks1 -S5 $FAST_PBKDF $LOOPDEV1 || fail
wipe $PWD1
check_hash $PWD1 $HASH1
echo $PWD1 | $REENC $LOOPDEV1 -q --decrypt || fail
check_hash_dev $LOOPDEV1 $HASH4

echo "[11] Detached header - adding encryption/reencryption/decryption"
prepare 8192
check_hash_dev $IMG $HASH4
echo $PWD1 | $REENC $LOOPDEV1 -q $FAST_PBKDF --header $IMG_HDR --new --type luks1
check_hash $PWD1 $HASH4 $IMG_HDR
echo $PWD1 | $REENC $LOOPDEV1 -q $FAST_PBKDF --header $IMG_HDR
check_hash $PWD1 $HASH4 $IMG_HDR
echo $PWD1 | $REENC $LOOPDEV1 -q --header $IMG_HDR --decrypt
check_hash_dev $IMG $HASH4
# existing header of zero size
cat /dev/null >$IMG_HDR
echo $PWD1 | $REENC $LOOPDEV1 -q $FAST_PBKDF --header $IMG_HDR --new --type luks1
check_hash $PWD1 $HASH4 $IMG_HDR
$CRYPTSETUP isLuks $LOOPDEV1 && fail
$CRYPTSETUP isLuks $IMG_HDR || fail

remove_mapping
exit 0