#!/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> - create active device"
	echo "        get-mode <name> - show verity mode of device"
	echo "        disable - disable dm-verity. reboot is needed"
}

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

		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=`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

	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
}

create()
{
	NAME=$1
	ROOTFS=$2

	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))

	((dmsetup_sectors=$block_count * $block_size / 512))
	((dmsetup_blocks=$block_count * $block_size / 4096))
	((dmsetup_hash_offset=$hash_offset / 4096 + 1))

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

	# The meta location is after the partition size. (After resizefs)
	part_block_size=`lsblk -rbno SIZE $ROOTFS`
	if [ $part_block_size -lt $hash_offset ]
	then
		echo "$ROOTFS does not have dm-verity meta data"
		exit 1
	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

			/sbin/dmsetup create $NAME -r --table "0 $dmsetup_sectors verity 1 $ROOTFS $ROOTFS 4096 4096 \
						  $dmsetup_blocks $dmsetup_hash_offset sha256 $root_hash $salt 1 ignore_zero_blocks"

			mount /dev/mapper/$NAME /sysroot
			exit 0
		elif [ "$enable" = b0b0b0b0 ]
		then
			echo "dm-verity is disabled (perform resize2fs)"
			/sbin/fsck -y $ROOTFS
			/sbin/resize2fs -f $ROOTFS
		else
			echo "$ROOTFS has the wrong enable flag"
		fi
	else
		echo "$ROOTFS is not dm-verity partition"
	fi

	exit 1
}

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=`/sbin/blkid -L rootfs`
	if [ -z "$ROOTFS" ]
	then
		ROOTFS=`/sbin/blkid -t PARTLABEL=rootfs -o device`
	fi

	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=root

	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=`/sbin/blkid -L rootfs`
	if [ -z "$ROOTFS" ]
	then
		ROOTFS=`/sbin/blkid -t PARTLABEL=rootfs -o device`
	fi

	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 -a ! -f /sbin/dmsetup ]
then
	exit 1
fi

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

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

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

	"disable")
		disable
		;;

	*)
		usage
		;;
esac

exit 0
