#!/bin/bash

[ -z "$CRYPTSETUP_PATH" ] && CRYPTSETUP_PATH=".."
CRYPTSETUP=$CRYPTSETUP_PATH/cryptsetup
DEV=""
DEV_STACKED="luks0xbabe"
DEV_NAME="dummyalign"
MNT_DIR="./mnt_luks"
PWD1="93R4P4pIqAH8"
PWD2="mymJeD8ivEhE"
FAST_PBKDF="--pbkdf-force-iterations 1000"

cleanup() {
	udevadm settle >/dev/null 2>&1
	if [ -d "$MNT_DIR" ] ; then
	    umount -f $MNT_DIR 2>/dev/null
	    rmdir $MNT_DIR 2>/dev/null
	fi
	[ -b /dev/mapper/$DEV_STACKED ] && dmsetup remove --retry $DEV_STACKED >/dev/null 2>&1
	[ -b /dev/mapper/$DEV_NAME ] && dmsetup remove --retry $DEV_NAME >/dev/null 2>&1
	# FIXME scsi_debug sometimes in-use here
	sleep 1
	rmmod scsi_debug 2>/dev/null
	sleep 1
}

fail()
{
	if [ -n "$1" ] ; then echo "FAIL $1" ; fi
	echo "FAILED backtrace:"
	while caller $frame; do ((frame++)); done
	cleanup
	exit 100
}

skip()
{
	echo "TEST SKIPPED: $1"
	cleanup
	exit 0
}

function dm_crypt_features()
{
	VER_STR=$(dmsetup targets | grep crypt | cut -f2 -dv)
	[ -z "$VER_STR" ] && fail "Failed to parse dm-crypt version."

	VER_MAJ=$(echo $VER_STR | cut -f 1 -d.)
	VER_MIN=$(echo $VER_STR | cut -f 2 -d.)
	VER_PTC=$(echo $VER_STR | cut -f 3 -d.)

	[ $VER_MAJ -lt 1 ] && return
	[ $VER_MAJ -gt 1 ] && {
		DM_PERF_CPU=1
		DM_SECTOR_SIZE=1
		return
	}

	[ $VER_MIN -lt 14 ] && return
	DM_PERF_CPU=1
	if [ $VER_MIN -ge 17 -o \( $VER_MIN -eq 14 -a $VER_PTC -ge 5 \) ]; then
		DM_SECTOR_SIZE=1
	fi
}

add_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
	DEV=$(grep -l -e scsi_debug /sys/block/*/device/model | cut -f4 -d /)

	if [ ! -e /sys/block/$DEV/alignment_offset ] ; then
		echo "This kernel seems to not support topology info, test skipped."
		cleanup
		exit 77
	fi

	DEV="/dev/$DEV"
	[ -b $DEV ] || fail "Cannot find $DEV."
}

format() # key_bits expected [forced]
{
	if [ -z "$3" ] ; then
		echo -n "Formatting using topology info ($1 bits key)..."
		echo $PWD1 | $CRYPTSETUP luksFormat --type luks1 $DEV -q $FAST_PBKDF -c aes-cbc-essiv:sha256 -s $1 || fail
	else
		echo -n "Formatting using forced sector alignment $3 ($1 bits key)..."
		echo $PWD1 | $CRYPTSETUP luksFormat --type luks1 $DEV -q $FAST_PBKDF -s $1 -c aes-cbc-essiv:sha256 --align-payload=$3 ||fail
	fi

	# check the device can be activated
	echo $PWD1 | $CRYPTSETUP luksOpen $DEV $DEV_NAME || fail
	$CRYPTSETUP close $DEV_NAME || fail

	ALIGN=$($CRYPTSETUP luksDump $DEV |grep "Payload offset" | sed -e s/.*\\t//)
	#echo "ALIGN = $ALIGN"

	[ -z "$ALIGN" ] && fail
	[ $ALIGN -ne $2 ] && fail "Expected alignment differs: expected $2 != detected $ALIGN"

	# test some operation, just in case
	echo -e "$PWD1\n$PWD2" | $CRYPTSETUP luksAddKey $DEV $FAST_PBKDF --key-slot 1
	[ $? -ne 0 ] && fail "Keyslot add failed."

	$CRYPTSETUP -q luksKillSlot $DEV 1
	[ $? -ne 0 ] && fail "Keyslot removal failed."

	echo "PASSED"
}

get_offsets()
{
	$CRYPTSETUP luksDump $DEV | grep "$1" | cut -s -d ':' -f 2 | sed  -e 's/\s//g' -e :a -e N -e 's/\n/:/g' -e 's/\s//g' -e ta
}

format_null()
{
	if [ $3 -eq 0 ] ; then
		echo -n "Formatting using topology info ($1 bits key) [slot 0"
		echo | $CRYPTSETUP luksFormat --type luks1 $DEV -q $FAST_PBKDF -c null -s $1 || fail
	else
		echo -n "Formatting using forced sector alignment $3 ($1 bits key) [slot 0"
		echo | $CRYPTSETUP luksFormat --type luks1 $DEV -q $FAST_PBKDF -c null -s $1 --align-payload=$3 || fail
	fi

	# check the device can be activated
	echo | $CRYPTSETUP luksOpen $DEV $DEV_NAME || fail
	$CRYPTSETUP close $DEV_NAME || fail

	POFF=$(get_offsets "Payload offset")
	[ -z "$POFF" ] && fail
	[ $POFF != $2 ] && fail "Expected data offset differs: expected $2 != detected $POFF"
	if [ -n "$4" ] ; then
		for j in 1 2 3 4 5 6 7 ; do
			echo -e "\n" | $CRYPTSETUP luksAddKey $DEV -q $FAST_PBKDF --key-slot $j -c null $PARAMS
			echo -n $j
			[ $? -ne 0 ] && fail
		done

		KOFF=$(get_offsets "Key material offset")
		[ -z "$KOFF" ] && fail
		[ $KOFF != $4 ] && fail "Expected keyslots offsets differ: expected $4 != detected $KOFF"
	fi

	echo "]...PASSED"
}

format_plain() # sector size
{
	echo -n "Formatting plain device (sector size $1)..."
	if [ -n "$DM_SECTOR_SIZE" ] ; then
		echo $PWD1 | $CRYPTSETUP open --type plain --hash sha256 --sector-size $1 $DEV $DEV_NAME || fail
		$CRYPTSETUP close $DEV_NAME || fail
		echo "PASSED"
	else
		echo "N/A"
	fi
}

format_plain_fail() # sector size
{
	echo -n "Formatting plain device (sector size $1, must fail)..."
	if [ -n "$DM_SECTOR_SIZE" ] ; then
		echo $PWD1 | $CRYPTSETUP open --type plain --hash sha256 --sector-size $1 $DEV $DEV_NAME >/dev/null 2>&1 && fail
		echo "PASSED"
	else
		echo "N/A"
	fi
}

if [ $(id -u) != 0 ]; then
	echo "WARNING: You must be root to run this test, test skipped."
	exit 77
fi

dm_crypt_features
modprobe --dry-run scsi_debug || exit 77
cleanup

echo "# Create desktop-class 4K drive"
echo "# (logical_block_size=512, physical_block_size=4096, alignment_offset=0)"
add_device dev_size_mb=16 sector_size=512 physblk_exp=3 num_tgts=1
format 256 4096
format 256 2056 8
format 128 2048
format 128 1032 8
format 256 8192 8192
format 128 8192 8192
cleanup

echo "# Create desktop-class 4K drive with misaligned opt-io (some bad USB enclosures)"
echo "# (logical_block_size=512, physical_block_size=4096, alignment_offset=0, opt-io=1025)"
add_device dev_size_mb=16 sector_size=512 physblk_exp=3 num_tgts=1 opt_blks=1025
format 256 4096
format 256 2056 8
format 128 2048
format 128 1032 8
format 256 8192 8192
format 128 8192 8192
cleanup

echo "# Create desktop-class 4K drive w/ 63-sector DOS partition compensation"
echo "# (logical_block_size=512, physical_block_size=4096, alignment_offset=3584)"
add_device dev_size_mb=16 sector_size=512 physblk_exp=3 lowest_aligned=7 num_tgts=1
format 256 4103
format 256 2056 8
format 128 2055
format 128 1032 8
cleanup

echo "# Create enterprise-class 4K drive"
echo "# (logical_block_size=4096, physical_block_size=4096, alignment_offset=0)"
add_device dev_size_mb=16 sector_size=4096 num_tgts=1 opt_blks=64
format 256 4096
format 256 2056 8
format 128 2048
format 128 1032 8
cleanup

echo "# Create classic 512B drive and stack dm-linear"
echo "# (logical_block_size=512, physical_block_size=512, alignment_offset=0)"
add_device dev_size_mb=16 sector_size=512 num_tgts=1
DEV2=$DEV
DEV=/dev/mapper/$DEV_STACKED
dmsetup create $DEV_STACKED --table "0 32768 linear $DEV2 0"
format 256 4096
format 256 2056 8
format 128 2048
format 128 1032 8
format 128 8192 8192
cleanup

echo "# Create classic 512B drive and stack dm-linear (plain mode)"
add_device dev_size_mb=16 sector_size=512 num_tgts=1
DEV2=$DEV
DEV=/dev/mapper/$DEV_STACKED
dmsetup create $DEV_STACKED --table "0 32768 linear $DEV2 0"
format_plain 512
format_plain 1024
format_plain 2048
format_plain 4096
format_plain_fail 1111
format_plain_fail 8192
echo "# Create classic 512B drive, unaligned to 4096 and stack dm-linear (plain mode)"
dmsetup remove --retry $DEV_STACKED >/dev/null 2>&1
dmsetup create $DEV_STACKED --table "0 32762 linear $DEV2 0"
format_plain 512
format_plain 1024
format_plain_fail 2048
format_plain_fail 4096
cleanup

echo "# Offset check: 512B sector drive"
add_device dev_size_mb=16 sector_size=512 num_tgts=1
#           |k| expO reqO expected slot offsets
format_null  64 2048    0 8:72:136:200:264:328:392:456
format_null  64  520    1
format_null  64  520    8
format_null  64  640  128
format_null  64 2048 2048
format_null 128 2048    0 8:136:264:392:520:648:776:904
format_null 128 1032    1
format_null 128 1032    8
format_null 128 1152  128
format_null 128 2048 2048
format_null 256 4096    0 8:264:520:776:1032:1288:1544:1800
format_null 256 2056    1
format_null 256 2056    8
format_null 256 2176  128
format_null 256 4096 2048
format_null 512 4096    0 8:512:1016:1520:2024:2528:3032:3536
format_null 512 4040    1
format_null 512 4040    8
format_null 512 4096  128
format_null 512 4096 2048
cleanup

echo "# Offset check: 4096B sector drive"
add_device dev_size_mb=16 sector_size=4096 num_tgts=1 opt_blks=64
format_null  64 2048    0 8:72:136:200:264:328:392:456
format_null  64  520    1
format_null  64  520    8
format_null  64  640  128
format_null  64 2048 2048
format_null 128 2048    0 8:136:264:392:520:648:776:904
format_null 128 1032    1
format_null 128 1032    8
format_null 128 1152  128
format_null 128 2048 2048
format_null 256 4096    0 8:264:520:776:1032:1288:1544:1800
format_null 256 2056    1
format_null 256 2056    8
format_null 256 2176  128
format_null 256 4096 2048
format_null 512 4096    0 8:512:1016:1520:2024:2528:3032:3536
format_null 512 4040    1
format_null 512 4040    8
format_null 512 4096  128
format_null 512 4096 2048
cleanup

echo "# Create enterprise-class 4K drive with fs and LUKS images."
# loop device here presents 512 block but images have 4k block
# cryptsetup should properly use 4k block on direct-io
add_device dev_size_mb=32 sector_size=4096 physblk_exp=0 num_tgts=1 opt_blks=64
for file in $(ls img_fs_*.img.xz) ; do
    echo "Format using fs image $file."
    xz -d -c $file | dd of=$DEV bs=1M 2>/dev/null || fail "bad image"
    [ ! -d $MNT_DIR ] && mkdir $MNT_DIR
    mount $DEV $MNT_DIR || skip "Mounting image is not available."
    echo $PWD1 | $CRYPTSETUP luksFormat --type luks1 --key-size 256 $FAST_PBKDF $MNT_DIR/luks.img || fail
    echo $PWD2 | $CRYPTSETUP luksFormat --type luks1 --key-size 256 $FAST_PBKDF $MNT_DIR/luks.img --header $MNT_DIR/luks_header.img || fail
    umount $MNT_DIR
done
cleanup