ddd

#!/bin/bash

set -e

DD_IF="if"
DD_OF="of"
DD_BS="bs" # takes precedence over ibs/obs
DD_IBS="ibs"
DD_OBS="obs"
DD_CNT="count"
DD_STS="status"
DD_SKIP="skip"
DD_SEEK="seek"

DDD_WBS="wbs" # write block size
DDD_WR="wr" # write retry
DDD_CMP="cmp" # write compare blocks

DD_STS_NONE="none"
DD_STS_NOXFER="noxfer"
DD_STS_PROGRESS="progress"

DDD_INPUT_SIZE="DDD_INPUT_SIZE"
DDD_OUTPUT_SIZE="DDD_OUTPUT_SIZE"

DDD_INPUT_BS="DDD_INPUT_BS"
DDD_OUTPUT_BS="DDD_OUTPUT_BS"
DDD_READ_START="DDD_READ_START" # read start byte offset
DDD_READ_SIZE="DDD_READ_SIZE" # read size byte length

DDD_WRITE_RETRY_CNT="DDD_WRITE_RETRY_CNT" # TODO: implement logic
DDD_CMP_CNT="DDD_CMP_CNT" # TODO: implement logic

DDD_WRITE_BS="DDD_WRITE_BS"
DDD_WRITE_START="DDD_WRITE_START"
DDD_WRITE_END="DDD_WRITE_END"

DDD_WRITE_SBLK="DDD_WRITE_SBLK" # start block number based on write_bs
DDD_WRITE_SOFF="DDD_WRITE_SOFF" # start block offset
DDD_WRITE_EBLK="DDD_WRITE_EBLK" # end block number based on write_bs
DDD_WRITE_EOFF="DDD_WRITE_EOFF" # end block offset

if [ -z "$DDD_WRITE_BS_DEFAULT_SHIFT" ]; then
	DDD_WRITE_BS_DEFAULT_SHIFT=22 # number of bits to shift to get write block size, 20 is 1MB, 21 is 2MB, 22 is 4MB
fi

# command line parameters
declare -A ddd_operands=()

# internal state parameters
declare -A ddd_internal=(
	[$DDD_INPUT_BS]=512 #assumed device sector size, TODO add sector size detection
	[$DDD_OUTPUT_BS]=512 #assumed device sector size, TODO add sector size detection
	[$DDD_WRITE_BS]=$((1 << $DDD_WRITE_BS_DEFAULT_SHIFT))
	[$DDD_WRITE_RETRY_CNT]=3
	[$DDD_CMP_CNT]=16
)

ddd_cmd="dd"

ddd_dump(){
	for key in ${!ddd_operands[@]}; do
		echo "$key: ${ddd_operands[$key]}" >&2
	done
	if [ ! -z "$1" ]; then
		for key in ${!ddd_internal[@]}; do
			local value=${ddd_internal[$key]}
			local value_unit=$(ddd_convertInt2BSS ${ddd_internal[$key]})
			if [ "$value" = "$value_unit" ]; then
				echo "$key:	$value" >&2
			else
				echo "$key:	$value (${value_unit})" >&2
			fi
		done
	fi
}

# run system command
ddd_system(){
	if ! local system_cmd=$(which "$1"); then
		echo "$FUNCNAME: required executable $1 missing" >&2
		return 1
	fi
	if [ ! -x "$system_cmd" ]; then
		echo "$FUNCNAME: no execute permission for $1" >&2
		return 1
	fi
	"$@"
}

# require dd if and of operands
ddd_checkOperands(){
	if [ -z "${ddd_operands[$DD_IF]}" ]; then
		echo "$FUNCNAME: required $DD_IF operand missing" >&2
		return 1
	fi
	if [ -z "${ddd_operands[$DD_OF]}" ]; then
		echo "$FUNCNAME: required $DD_OF operand missing" >&2
		return 1
	fi
}

# set ibs based on bs or ibs operand
ddd_setInternalInputBS(){
	if [ ! -z "${ddd_operands[$DD_BS]}" ]; then
		ddd_internal[$DDD_INPUT_BS]=${ddd_operands[$DD_BS]}
	elif [ ! -z "${ddd_operands[$DD_IBS]}" ]; then
		ddd_internal[$DDD_INPUT_BS]=${ddd_operands[$DD_IBS]}
	fi
}

# check if operands and set internal parameters
ddd_checkRead(){
	ddd_setInternalInputBS
	
	if [ -b "${ddd_operands[$DD_IF]}" ]; then
		local input_size_B=$(ddd_system blockdev --getsize64 "${ddd_operands[$DD_IF]}")
	elif [ -f "${ddd_operands[$DD_IF]}" ]; then
		local input_size_B=$(stat -c %s "${ddd_operands[$DD_IF]}")
	elif [ -c "${ddd_operands[$DD_IF]}" ]; then
		echo "$FUNCNAME: $DD_IF operand is a character device" >&2
		return 1
	else
		echo "$FUNCNAME: $DD_IF operand is not a block device or file" >&2
		return 1
	fi
	ddd_internal[$DDD_INPUT_SIZE]=$input_size_B

	local ibs=${ddd_internal[$DDD_INPUT_BS]}

	local skip_B=0 # bytes to skip
	if [ ! -z "${ddd_operands[$DD_SKIP]}" ]; then
		local skip_B=$((ibs * ${ddd_operands[$DD_SKIP]}))
	fi
	ddd_internal[$DDD_READ_START]=$skip_B

	local size_B=0
	if [ ! -z "${ddd_operands[$DD_CNT]}" ]; then
		local size_B=$((ibs * ${ddd_operands[$DD_CNT]}))
	else
		local size_B=$((input_size_B - $skip_B))
	fi
	ddd_internal[$DDD_READ_SIZE]=$size_B
	local end_B=$((skip_B + size_B))
	if [ "$end_B" -gt "$input_size_B" ]; then
		echo "$FUNCNAME: read to $end_B exceeds input size $input_size_B by $(($end_B - $input_size_B)) bytes" >&2
		return 1
	fi
}

# set obs based on bs or obs operand
ddd_setInternalOutputBS(){
	local ddd_write_bs=${ddd_internal[$DDD_WRITE_BS]}
	if [ ! -z "${ddd_operands[$DD_BS]}" ]; then
		local ddd_output_bs=${ddd_operands[$DD_BS]}
	elif [ ! -z "${ddd_operands[$DD_OBS]}" ]; then
		local ddd_output_bs=${ddd_operands[$DD_OBS]}
	fi
	if [ ! -z "$ddd_output_bs" ]; then
		if [ $ddd_output_bs -gt $ddd_write_bs ]; then
			if [ $(($ddd_output_bs % $ddd_write_bs)) -ne 0 ]; then
				echo "$FUNCNAME: output block size must be a multiple of 2."
				return 1
			fi
			ddd_internal[$DDD_WRITE_BS]=$ddd_output_bs
		fi
		ddd_internal[$DDD_OUTPUT_BS]=$ddd_output_bs
	fi
}

# check of operands and set internal parameters
ddd_checkWrite(){
	ddd_setInternalOutputBS
	local output_size_B=-1	
	if [ -b "${ddd_operands[$DD_OF]}" ]; then
		local output_size_B=$(ddd_system blockdev --getsize64 "${ddd_operands[$DD_OF]}")
	fi
	
	local obs=${ddd_internal[$DDD_OUTPUT_BS]}

	local seek_B=0
	if [ ! -z "${ddd_operands[$DD_SEEK]}" ]; then
		local seek_B=$((obs * ${ddd_operands[$DD_SEEK]}))
	fi
	ddd_internal[$DDD_WRITE_START]=$seek_B
	
	local end_B=$((seek_B + ${ddd_internal[$DDD_READ_SIZE]}))
	ddd_internal[$DDD_WRITE_END]=$end_B

	if [ "$output_size_B" -ge 0 ] && [ "$end_B" -gt "$output_size_B" ]; then
		echo "$FUNCNAME: write to $end exceeds output size $output_size by $((end_B - $output_size_B))" >&2
		return 1
	fi
	ddd_internal[$DDD_WRITE_SOFF]=$((${ddd_internal[$DDD_WRITE_START]} % ${ddd_internal[$DDD_WRITE_BS]}))
	ddd_internal[$DDD_WRITE_SBLK]=$((${ddd_internal[$DDD_WRITE_START]} / ${ddd_internal[$DDD_WRITE_BS]}))
	ddd_internal[$DDD_WRITE_EOFF]=$((${ddd_internal[$DDD_WRITE_END]} % ${ddd_internal[$DDD_WRITE_BS]}))
	ddd_internal[$DDD_WRITE_EBLK]=$(((${ddd_internal[$DDD_WRITE_END]} - 1) / ${ddd_internal[$DDD_WRITE_BS]}))
}

# convert int to bs with suffix
ddd_convertInt2BSS(){
	local digit="$1"
	local base=1024
	local exponent=0
	while [ "$digit" -ge "$base" ]; do
		if [ "$((digit % base))" -eq 0 ]; then
			local digit=$((digit >> 10))
			local exponent=$((exponent+1))
		else
			break
		fi
	done
	if [ "$exponent" -eq 0 ]; then
		local base=1000
		while [ "$digit" -ge "$base" ]; do
			if [ "$((digit % base))" -eq 0 ]; then
				local digit=$((digit / base))
				local exponent=$((exponent+1))
			else
				break
			fi
		done
	fi
	case "$exponent" in
		0)
			local suffix=
			;;
		1)
			local suffix=K
			;;
		2)
			local suffix=M
			;;
		3)
			local suffix=G
			;;
		4)
			local suffix=T
			;;
		5)
			local suffix=P
			;;
		6)
			local suffix=E
			;;
		7)
			local suffix=Z
			;;
		8)
			local suffix=Y
			;;
		*)
			echo "$FUNCNAME: suffix is not supported" >&2
			return 1
			;;
	esac
	if [ "$base" -eq 1000 ] && [ "$exponent" -gt 0 ]; then
		local suffix="${suffix}B"
	fi
	echo -n "$digit$suffix"
}

# convert bs with suffix to integer
ddd_convertBSS2Int(){
	#empty bs is prefiltered by main
	#just a number, return number without leading 0s
	if [ -z "${1//[0-9]}" ]; then
		echo "${1##0*}"
		return
	fi

	local digit=$(echo "$1" | sed "s/[^0-9]*\$//")
	local suffix=$(echo "$1" | sed "s/^[0-9]*//")

	if [ "$1" != "$digit$suffix" ]; then
		echo "$FUNCNAME: $DD_BS operand is invalid" >&2
		return 1
	fi

	if [ -z "$digit" ]; then
		local digit=1
	fi

	if [ -z "$suffix" ]; then
		echo -n "$digit"
		return
	fi

	if [ "${#suffix}" -gt 3 ]; then
		echo "$FUNCNAME: $DD_BS operand suffix is unsupported" >&2
		return 1
	fi
	local base=1
	local exponent=1
	case "$suffix" in
		"c")
			:
			;;
		"w")
			local digit=$((digit<<1))
			;;
		"b")
			local digit=$((digit<<9))
			;;
		*)
			local base=1024
			if [ "${#suffix}" = 2 ]; then
				if [ "${suffix:1:1}" != "B" ]; then
					echo "$FUNCNAME: $DD_BS operand suffix is unsupported" >&2
					return 1
				fi
				local base=1000
			elif [ "${#suffix}" = 3 ]; then
				if [ "${suffix:1:2}" != "iB" ]; then
					echo "$FUNCNAME: $DD_BS operand suffix is unsupported" >&2
					return 1
				fi
			fi
			local suffix="${suffix:0:1}"
			case "${suffix,,}" in
				"k")
					local exponent=1
					;;
				"m")
					local exponent=2
					;;
				"g")
					local exponent=3
					;;
				"t")
					local exponent=4
					;;
				"p")
					local exponent=5
					;;
				"e")
					local exponent=6
					;;
				"z")
					local exponent=7
					;;
				"y")
					local exponent=8
					;;
				*)
					echo "$FUNCNAME: $DD_BS operand suffix is unsupported" >&2
					return 1
			esac
			;;
	esac
	echo -n "$((digit*(base**exponent)))"
}

# write a full block, 1 skip, 2 count, 3 seek
ddd_syncBlock(){
	#set -x
	if cmp -s -n $ddd_obs	<(dd if="$ddd_if" ibs=$ddd_ibs skip=$1 count=$2 status=none) \
				<(dd if="$ddd_of" ibs=$ddd_obs skip=$3 count=1 status=none iflag=nocache); then
		ddd_blocks_skip=$((++ddd_blocks_skip))
		echo -n "SKIP"
		return
	fi
	ddd_blocks_sync=$((++ddd_blocks_sync))
	echo -n "SYNC"
	if ! dd if="$ddd_if" ibs=$ddd_ibs skip=$1 count=$2 of="$ddd_of" obs=$ddd_obs seek=$3 status=none oflag=sync conv=notrunc; then
		ddd_blocks_sync_fail=$((++ddd_blocks_sync_fail))
		echo "SYNC FAIL"
		echo "$FUNCNAME: failed write at write block $ddd_block" >&2
		return 2
	fi
	if ! cmp -s -n $ddd_obs	<(dd if="$ddd_if" ibs=$ddd_ibs skip=$1 count=$2 status=none) \
				<(dd if="$ddd_of" ibs=$ddd_obs skip=$3 count=1 status=none iflag=nocache); then
		ddd_blocks_cmp1_fail=$((++ddd_blocks_cmp1_fail))
		echo "CMP FAIL"
		echo "$FUNCNAME: mismatch at block $ddd_block" >&2
		return 1
	fi
	#set +x
}

# write a full block, 1 skip, 2 seek
ddd_syncPartialBlock(){
	#set -x
	if cmp -s -n $ddd_bs_single	<(dd if="$ddd_if" ibs=$ddd_bs_single skip=$1 count=1 status=none) \
					<(dd if="$ddd_of" ibs=$ddd_bs_single skip=$2 count=1 status=none iflag=nocache); then
		ddd_blocks_skip=$((++ddd_blocks_skip))
		echo -n "SKIP"
		return
	fi
	ddd_blocks_sync=$((++ddd_blocks_sync))
	echo -n "SYNC"
	if ! dd if="$ddd_if" ibs=$ddd_bs_single skip=$1 count=1 of="$ddd_of" obs=$ddd_bs_single seek=$2 status=none oflag=sync conv=notrunc; then
		ddd_blocks_sync_fail=$((++ddd_blocks_sync_fail))
		echo "SYNC FAIL"
		echo "$FUNCNAME: failed write at write block $ddd_block" >&2
		return 2
	fi
	if ! cmp -s -n $ddd_bs_single	<(dd if="$ddd_if" ibs=$ddd_bs_single skip=$1 count=1 status=none) \
					<(dd if="$ddd_of" ibs=$ddd_bs_single skip=$2 count=1 status=none iflag=nocache); then
		ddd_blocks_cmp1_fail=$((++ddd_blocks_cmp1_fail))
		echo "CMP FAIL"
		echo "$FUNCNAME: mismatch at block $ddd_block" >&2
		return 1
	fi
	#set +x
}

ddd_sync(){
	# TODO: write retry
	# TODO: compare blocks
	# TODO: check count * ibs = obs
	
	#flatten variables
	ddd_if=${ddd_operands[$DD_IF]}
	ddd_of=${ddd_operands[$DD_OF]}

	ddd_ibs=${ddd_internal[$DDD_INPUT_BS]}
	ddd_obs=${ddd_internal[$DDD_WRITE_BS]}
	
	ddd_read_start=${ddd_internal[$DDD_READ_START]}
	ddd_write_start=${ddd_internal[$DDD_WRITE_START]}
	
	ddd_write_sblk=${ddd_internal[$DDD_WRITE_SBLK]}
	ddd_write_eblk=${ddd_internal[$DDD_WRITE_EBLK]}

	ddd_block=0
	ddd_blocks=$((ddd_write_eblk - ddd_write_sblk + 1))

	ddd_blocks_skip=0
	ddd_blocks_sync=0
	ddd_blocks_sync_fail=0
	ddd_blocks_cmp1_fail=0

	#set -x
	for i in $(seq $ddd_write_sblk $ddd_write_eblk); do
		ddd_block=$((ddd_block + 1))
		echo -en "\rddd: block $ddd_block of $ddd_blocks: "
		if [ $i -ne $ddd_write_sblk ] && [ $i -ne $ddd_write_eblk ]; then
			local skip=$((ddd_read_start + ddd_bs_single + (i - ddd_write_sblk - 1) * ddd_obs))
			local count=$((ddd_obs / ddd_ibs))
			ddd_syncBlock ${skip}B ${count} ${i} # TODO: implement retry logic
		else
			local seek=$((ddd_write_start + (i - ddd_write_sblk) * ddd_obs))
			if [ $i -eq $ddd_write_sblk ] && [ $i -eq $ddd_write_eblk ]; then
				local skip=$ddd_read_start
				local seek=$ddd_write_start
				ddd_bs_single=${ddd_internal[$DDD_READ_SIZE]}
			elif [ $i -eq $ddd_write_sblk ]; then
				local skip=$ddd_read_start
				local seek=$ddd_write_start
				ddd_bs_single=$((ddd_obs - ${ddd_internal[$DDD_WRITE_SOFF]}))
			elif [ $i -eq $ddd_write_eblk ]; then
				local skip=$((ddd_read_start + ddd_bs_single + (i - ddd_write_sblk - 1) * ddd_obs))
				local seek=$((ddd_write_start + ddd_bs_single + (i - ddd_write_sblk - 1) * ddd_obs))
				ddd_bs_single=${ddd_internal[$DDD_WRITE_EOFF]}
				if [ $ddd_bs_single -eq 0 ]; then
					ddd_bs_single=$ddd_obs
				fi
			fi
			ddd_syncPartialBlock ${skip}B ${seek}B
		fi
	done
	#set +x
	echo -e "\rSKIP: $ddd_blocks_skip SYNC: $ddd_blocks_sync SYNC_FAIL: $ddd_blocks_sync_fail CMP1_FAIL: $ddd_blocks_cmp1_fail"
}

ddd_verify(){
	cmp -s -n "${ddd_internal[$DDD_READ_SIZE]}" <() <()
}

ddd_debug(){
	if [ ! -z "$DDD_DEBUG" ]; then
		echo "$FUNCNAME: $@" >&2
	fi
}

ddd_main(){
	while [ ! -z "$1" ]; do
		local operand="${1%%=*}"
		local operand_value=${1#$operand=}
		ddd_debug "$FUNCNAME" "$operand" "$operand_value"
		case "$operand" in
			$DD_IF)
				if [ ! -e "$operand_value" ]; then
					echo "$FUNCNAME: $operand operand value does not exist" >&2
					return 1
				fi
				if [ ! -r "$operand_value" ]; then
					echo "$FUNCNAME: $operand operand value is not readable by current user" >&2
					return 1
				fi
				ddd_operands[$operand]="$operand_value"
				;;
			$DD_OF)
				if [ ! -e "$operand_value" ]; then
					if [ ! -w "$(dirname "$operand_value")" ]; then # TODO: is this correct way to determining parent directory writable?
						echo "$FUNCNAME: $operand operand directory is not writable by current user" >&2
						return 1
					fi
				elif [ ! -w "$operand_value" ]; then
					echo "$FUNCNAME: $operand operand value is not writable by current user" >&2
					return 1
				fi
				ddd_operands[$operand]="$operand_value"
				;;
			$DD_BS)
				if [ ! -z "${ddd_operands[$DD_IBS]}" ]; then
					echo "$FUNCNAME: $operand operand priority over $DD_IBS" >&2
					unset ddd_operands[$DD_IBS]
				fi
				if [ ! -z "${ddd_operands[$DD_OBS]}" ]; then
					echo "$FUNCNAME: $operand operand priority over $DD_OBS" >&2
					unset ddd_operands[$DD_OBS]
				fi
				#convert block size with suffix to integer
				ddd_operands[$operand]="$(ddd_convertBSS2Int "$operand_value")"
				;;
			$DD_IBS)
				if [ ! -z "${ddd_operands[$DD_BS]}" ]; then
					echo "$FUNCNAME: $DD_BS operand priority over $operand" >&2
				fi
				echo "$FUNCNAME: $operand applies to skip only" >&2
				#convert block size with suffix to integer
				ddd_operands[$operand]="$(ddd_convertBSS2Int "$operand_value")"
				
				;;
			$DD_OBS)
				if [ ! -z "${ddd_operands[$DD_BS]}" ]; then
					echo "$FUNCNAME: $DD_BS operand priority over $operand" >&2
				fi
				echo "$FUNCNAME: $operand applies to seek only, new wbs operand determines write block size" >&2
				#convert block size with suffix to integer
				ddd_operands[$operand]="$(ddd_convertBSS2Int "$operand_value")"
				;;
			$DD_CNT)
				if [ ! -z "${operand_value//[0-9]}" ]; then
					echo "$FUNCNAME: $operand operand is not a valid integer" >&2
					return 1
				fi
				#remove preceding zeros
				ddd_operands[$operand]="${operand_value##0*}"
				;;
			$DD_STS)
				case "$operand_value" in
					$DD_STS_NONE)
						:
						;;
					$DD_STS_NOXFER)
						:
						;;
					$DD_STS_PROGRESS)
						:
						;;
					*)
						echo "$FUNCNAME: $operand operand is not valid" >&2
						return 1
						;;
				esac
				ddd_operands[$operand]="$operand_value"
				;;
			$DD_SKIP)
				if [ ! -z "$(echo "${operand_value: -1}" | sed "s/[0-9B]//")" ]; then
					echo "$FUNCNAME: $operand operand must be an integer representing a multliple of the block size or an byte offset" >&2
					return 1
				elif [ -z "$operand_value" ]; then
					echo "$FUNCNAME: $operand value must not be empty" >&2
					return 1
				fi
				ddd_operands[$operand]="$operand_value"
				;;
			$DD_SEEK)
				if [ ! -z "$(echo "${operand_value: -1}" | sed "s/[0-9B]//")" ]; then
					echo "$FUNCNAME: $operand operand must be an integer representing a multiple of the block size or an byte offset" >&2
					return 1
				elif [ -z "$operand_value" ]; then
					echo "$FUNCNAME: $operand value must not be empty" >&2
					return 1
				fi
				ddd_operands[$operand]="$operand_value"
				;;
			$DDD_WBS)
				#convert block size with suffix to integer
				echo "$FUNCNAME: overriding write block size $(ddd_convertInt2BSS ${ddd_internal[$DDD_WRITE_BS]}) with $operand_value" >&2
				ddd_internal[$DDD_WRITE_BS]="$(ddd_convertBSS2Int "$operand_value")"
				;;
			*)
				echo "$FUNCNAME: unrecognized operand $operand." >&2
				exit 1
				;;
		esac
		shift
	done

	ddd_checkOperands

	for operand in "${!ddd_operands[@]}"; do
		if [ -z "${ddd_operands[$operand]}" ]; then
			ddd_cmd+=" $operand"
		elif [ "$operand" = $DD_BS ] || [ "$operand" = $DD_IBS ] || [ "$operand" = $DD_OBS ]; then
			ddd_cmd+=" $operand=$(ddd_convertInt2BSS ${ddd_operands[$operand]})"
		else
			ddd_cmd+=" $operand=${ddd_operands[$operand]}"
		fi
	done
	echo $ddd_cmd >&2

	ddd_checkRead
	ddd_checkWrite

	ddd_dump 1
	
	#set -x
	if [ ! -e "${ddd_operands[$DD_OF]}" ]; then
		$ddd_cmd conv=sparse status=progress
	fi
	#set +x
	
	ddd_sync

	#ddd_verifyWrite
}

ddd_main "$@"