#!/bin/sh

PATH=/bin:/usr/bin:/sbin:/usr/sbin

usage()
{
	echo "Usage: verityctl <action> <action-specific>"
	echo ""
	echo "Action commands:"
	echo "        format   <image> - create hashtable"
	echo "        create   <name> <device> <mount point> - create active device"
	echo "        close    <name> - close dm-verity device"
	echo "        get-mode <name> - show verity mode of device"
	echo "        disable - disable dm-verity. reboot is needed"
}

get_rootfs() {
	# get partition ab
	P_SLOT=$([[ $(</proc/cmdline) =~ partition_ab=([ab]) ]]; echo ${BASH_REMATCH[1]})
	P_SUFFIX=""

	if [ "${P_SLOT}" != "" ]; then
		P_SUFFIX="_${P_SLOT}"
	fi

	ROOTFS=`/sbin/blkid -t PARTLABEL=rootfs${P_SUFFIX} -o device -l`
	if [ -z "$ROOTFS" ]
	then
		ROOTFS=`/sbin/blkid -t PARTLABEL=rootfs -o device -l`
	fi
	if [ -z "$ROOTFS" ]
	then
		ROOTFS=`/sbin/blkid -t LABEL=contain-rootfs -o device -l`
	fi

	echo "$ROOTFS"
}

format()
{
	IMG_FILE=$1

	if [ -f $IMG_FILE ] || [ -b $IMG_FILE ]
	then
		echo "Run verityctl format $IMG_FILE"
	else
		echo "$IMG_FILE does not exist"
		exit 1
	fi

	# Block device such as /dev/mmcblk0p1
	if [ -b $IMG_FILE ]
	then
		IMG_PATH=/tmp

		# Only support ext4 filesystem
		FS_TYPE=`/bin/lsblk -n -o FSTYPE $IMG_FILE`
		if [ x"$FS_TYPE" != x"ext4" ]
		then
			echo "$IMG_FILE is not ext4 filesystem"
			exit 0
		fi

		block_count=`/sbin/tune2fs -l $IMG_FILE | grep "Block count" | gawk '{print $3}'`
		block_size=`/sbin/tune2fs -l $IMG_FILE | grep "Block size" | gawk '{print $3}'`

		((meta_offset=$block_count * $block_size))
		part_block_size=`/bin/lsblk -rbno SIZE $IMG_FILE`

		echo "partition size: $part_block_size"
		echo "meta data offset: $meta_offset"

		if [ $part_block_size -le $meta_offset ]
		then
			echo "$IMG_FILE does not have dm-verity meta data"
			exit 0
		fi

		OPTS="--data-blocks=$block_count"
	else
		IMG_PATH=`dirname $IMG_FILE`
		OPTS=""
	fi

	/sbin/veritysetup format $OPTS $IMG_FILE $IMG_PATH/hash_data | tee $IMG_PATH/verity_format_output.txt

	root_hash=`grep "Root hash" $IMG_PATH/verity_format_output.txt | gawk '{print $3}'`
	salt=`grep "Salt" $IMG_PATH/verity_format_output.txt | gawk '{print $2}'`

	dd if=/dev/zero of=$IMG_PATH/meta_data bs=32768 count=1 2> /dev/null
	echo "dm-verity0" | dd of=/$IMG_PATH/meta_data bs=1 seek=0  conv=notrunc 2> /dev/null
	echo "b1b1b1b1"   | dd of=/$IMG_PATH/meta_data bs=1 seek=16 conv=notrunc 2> /dev/null
	echo $root_hash   | dd of=/$IMG_PATH/meta_data bs=1 seek=32 conv=notrunc 2> /dev/null
	echo $salt        | dd of=/$IMG_PATH/meta_data bs=1 seek=96 conv=notrunc 2> /dev/null

	# root_hash and salt signing - use of signing algorithm suitable for each product
	############################################################################################
	# echo ${root_hash}${salt} > $IMG_PATH/root_hash_salt.txt
	# /usr/bin/openssl dgst -sha512 -sign /etc/ssl/certs/verifyboot_private.pem -out $IMG_PATH/root_hash_salt.sha512 $IMG_PATH/root_hash_salt.txt
	# dd if=$IMG_PATH/root_hash_salt.sha512 of=/$IMG_PATH/meta_data bs=1 seek=256 conv=notrunc 2> /dev/null
	# rm -f $IMG_PATH/root_hash_salt.txt
	# rm -f $IMG_PATH/root_hash_salt.sha512
	############################################################################################

	if [ -b $IMG_FILE ]
	then
		((hash_offset=$meta_offset + 32768))

		hash_size=`ls -s --block-size=1 $IMG_PATH/hash_data | gawk '{print $1}'`
		((hash_end_offset=$hash_offset + $hash_size))

		echo "hash end offset: $hash_end_offset"

		if [ $part_block_size -lt $hash_end_offset ]
		then
			echo "$IMG_FILE does not have enough space for hash data"
			exit 1
		fi

		dd if=$IMG_PATH/meta_data of=$IMG_FILE bs=1 seek=$meta_offset 2> /dev/null
		dd if=$IMG_PATH/hash_data of=$IMG_FILE bs=1 seek=$hash_offset 2> /dev/null
	else
		cat $IMG_PATH/meta_data $IMG_PATH/hash_data >> $IMG_FILE
	fi

	rm -f $IMG_PATH/hash_data
	rm -f $IMG_PATH/meta_data
	rm -f $IMG_PATH/verity_format_output.txt
}

################################################################################################
# Return values
# 0 : normal boot
# 1 : verity boot is enabled
# 2 : verity boot is enabled but corrupted
# 3 : verify boot is disabled and should perform resizefs
################################################################################################
create()
{
	NAME=$1
	ROOTFS=$2
	MOUNT_POINT=$3

	block_count=`/sbin/tune2fs -l $ROOTFS | grep "Block count" | gawk '{print $3}'`
	block_size=`/sbin/tune2fs -l $ROOTFS | grep "Block size" | gawk '{print $3}'`

	((meta_offset=$block_count * $block_size))
	((hash_offset=$meta_offset + 32768))

	((meta_verity_offset=$meta_offset))
	((meta_enable_offset=$meta_offset + 16))
	((meta_root_hash_offset=$meta_offset + 32))
	((meta_salt_offset=$meta_offset + 96))
	((meta_sign_offset=$meta_offset + 256))

	# The meta location is after the partition size. (After resizefs)
	part_block_size=`/bin/lsblk -rbndo SIZE $ROOTFS`
	if [ $part_block_size -lt $hash_offset ]
	then
		echo "$ROOTFS does not have dm-verity meta data"
		exit 0
	fi

	verity=`dd if=$ROOTFS bs=1 skip=$meta_verity_offset count=10 2> /dev/null`
	enable=`dd if=$ROOTFS bs=1 skip=$meta_enable_offset count=8 2> /dev/null`
	root_hash=`dd if=$ROOTFS bs=1 skip=$meta_root_hash_offset count=64 2> /dev/null`
	salt=`dd if=$ROOTFS bs=1 skip=$meta_salt_offset count=64 2> /dev/null`

	if [ "$verity" = dm-verity0 ]
	then
		if [ "$enable" = b1b1b1b1 ]
		then
			/sbin/veritysetup --hash-offset=$hash_offset dump $ROOTFS  # only show dm-verity information

			# root_hash and salt signing - use signing algorithm suitable for each product
			############################################################################
			# echo ${root_hash}${salt} > /tmp/root_hash_salt.txt
			# dd if=$ROOTFS of=/tmp/root_hash_salt.sha512 bs=1 skip=$meta_sign_offset count=128 2> /dev/null
			# /usr/bin/openssl dgst -sha512 -verify /etc/ssl/certs/verifyboot_public.pem -signature /tmp/root_hash_salt.sha512 /tmp/root_hash_salt.txt
			# signing_result=$?
			# /bin/rm -f /tmp/root_hash_salt.txt
			# /bin/rm -f /tmp/root_hash_salt.sha512
			# if [ $signing_result != 0 ]; then exit 2; fi
			############################################################################
			if [ -x /usr/bin/dmverity-rootfs-verify-hash.sh ]; then
				/usr/bin/dmverity-rootfs-verify-hash.sh ${ROOTFS} ${root_hash} ${salt} ${meta_sign_offset}
				VERIFY_RESULT=$?
				if [ ${VERIFY_RESULT} -ne 0 ]; then
					echo "dm-verity root hash verification for ${ROOTFS} failed: ${VERIFY_RESULT}"
					exit 2;
				fi
				echo "dm-verity root hash verification for ${ROOTFS} succeeded"
			fi

			# replace dmsetup to veritysetup
			# veritysetup reads super block to retrieve parameters:
			# 	data block size, hash block size, number of data block, hashing algorithm, salt
			/sbin/veritysetup --hash-offset=$hash_offset open $ROOTFS $NAME $ROOTFS $root_hash --ignore-zero-blocks --ignore-corruption
			VERITYSETUP_RESULT=$?
			if [ ${VERITYSETUP_RESULT} -ne 0 ]; then
				echo "veritysetup failed: ${VERITYSETUP_RESULT}"
				exit 2;
			fi

			mount /dev/mapper/$NAME "$MOUNT_POINT"
			exit 1
		elif [ "$enable" = b0b0b0b0 ]
		then
			exit 3
		else
			echo "$ROOTFS has the wrong enable flag"
		fi
	else
		echo "$ROOTFS is not dm-verity partition"
	fi

	exit 0
}

verity_close() {
	NAME=$1

	/sbin/veritysetup status $NAME
	if [ $? -ne 0 ]; then
		echo "$NAME is not a dm-verity name"
		exit 1
	fi

	/sbin/veritysetup close $NAME

	if [ $? -eq 0 ]; then
		echo "veritysetup close $NAME: succeeded"
	else
		echo "veritysetup close $NAME: failed($VERITYSETUP_RESULT)"
	fi

	exit 0
}

get_mode()
{
	NAME=$1

	VERITY_BOOT=`/sbin/veritysetup status $NAME | grep "is active and is in use"`
	if [ -z "$VERITY_BOOT" ]
	then
		echo "dm-verity is disabled (Normal boot)"
		exit 0
	fi

	ROOTFS=$(get_rootfs)

	block_count=`/sbin/tune2fs -l $ROOTFS | grep "Block count" | gawk '{print $3}'`
	block_size=`/sbin/tune2fs -l $ROOTFS | grep "Block size" | gawk '{print $3}'`

	((meta_offset=$block_count * $block_size))
	((hash_offset=$meta_offset + 32768))

	((meta_verity_offset=$meta_offset))
	((meta_enable_offset=$meta_offset + 16))
	((meta_root_hash_offset=$meta_offset + 32))
	((meta_salt_offset=$meta_offset + 96))

	verity=`dd if=$ROOTFS bs=1 skip=$meta_verity_offset count=10 2> /dev/null`
	enable=`dd if=$ROOTFS bs=1 skip=$meta_enable_offset count=8 2> /dev/null`
	root_hash=`dd if=$ROOTFS bs=1 skip=$meta_root_hash_offset count=64 2> /dev/null`
	salt=`dd if=$ROOTFS bs=1 skip=$meta_salt_offset count=64 2> /dev/null`

	if [ "$verity" = dm-verity0 ]
	then
		if [ "$enable" = b1b1b1b1 ]
		then
			echo "dm-verity is enabled"
		elif [ "$enable" = b0b0b0b0 ]
		then
			echo "dm-verity is disabled"
		else
			echo "dm-verity is disabled (Unknown enable mode)"
		fi
	else
		echo "dm-verity is disabled (Bad magic number)"
	fi
}

disable()
{
	NAME=rootfs

	VERITY_BOOT=`/sbin/veritysetup status $NAME | grep "is active and is in use"`
	if [ -z "$VERITY_BOOT" ]
	then
		echo "dm-verity is disabled (Normal boot)"
		exit 0
	fi

	ROOTFS=$(get_rootfs)

	block_count=`/sbin/tune2fs -l $ROOTFS | grep "Block count" | gawk '{print $3}'`
	block_size=`/sbin/tune2fs -l $ROOTFS | grep "Block size" | gawk '{print $3}'`

	((meta_offset=$block_count * $block_size))
	((hash_offset=$meta_offset + 32768))

	((meta_verity_offset=$meta_offset))
	((meta_enable_offset=$meta_offset + 16))
	((meta_root_hash_offset=$meta_offset + 32))
	((meta_salt_offset=$meta_offset + 96))

	verity=`dd if=$ROOTFS bs=1 skip=$meta_verity_offset count=10 2> /dev/null`
	enable=`dd if=$ROOTFS bs=1 skip=$meta_enable_offset count=8 2> /dev/null`
	root_hash=`dd if=$ROOTFS bs=1 skip=$meta_root_hash_offset count=64 2> /dev/null`
	salt=`dd if=$ROOTFS bs=1 skip=$meta_salt_offset count=64 2> /dev/null`

	if [ "$verity" = dm-verity0 ]
	then
		if [ "$enable" = b1b1b1b1 ]
		then
			echo "dm-verity is disabled"
			echo "b0b0b0b0" | dd of=$ROOTFS bs=1 seek=$meta_enable_offset count=8 conv=notrunc 2> /dev/null
		elif [ "$enable" = b0b0b0b0 ]
		then
			echo "dm-verity is already disabled"
		else
			echo "dm-verity is already disabled (Unknown enable mode)"
		fi
	else
		echo "dm-verity is already disabled (Bad magic number)"
	fi
}

# if /sbin/veritysetup does not exist, ignore dm-verity.
if [ ! -f /sbin/veritysetup ]
then
	exit 1
fi

case $1 in
	"format")
		if [ $# -ne 2 ]; then usage; exit 1; fi
		format $2
		;;

	"create")
		if [ $# -ne 4 ]; then usage; exit 1; fi
		create $2 $3 $4
		;;

	"close")
		if [ $# -ne 2 ]; then usage; exit 1; fi
		verity_close $2
		;;

	"get-mode")
		if [ $# -ne 2 ]; then usage; exit 1; fi
		get_mode $2
		;;

	"disable")
		disable
		;;

	*)
		usage
		;;
esac

exit 0
