#!/bin/sh

set -e

PRELOAD_MAGIC="55425348"
HEADER_METADATA_LENGHT=64

# Default values
SIGNATURE_SPECS_SIGNALGO=rsa2048
SIGNATURE_SPECS_PADDING=pss
SIGNATURE_SPECS_HASHALGO=sha256
SIGNATURE_SPECS_PUBKEY="/security/SW_DESCRIPTION_SIGN_PUBKEY_FILE.pem"
PRELOAD_SIGNATURE_MANDATORY=true


USAGE="Usage: $0 MODE FILE OPTIONS

This is a tool to check and get information from U-boot FIT images.

MODE:
    MODE is the U-boot FIT tool selector. It should be one of:
        - info: retrieve FIT properties.

FILE:
    File to be used as FIT. The syntax defines the file type:
        - simple file: direct path to file [/my/file]
        - UBI volume: UBI device and volume name [ubi0:rootfs]

OPTIONS:

    Pre-load header related options:
        -m              Make pre-load header signature mandatory. If this flag is set,
                        not finding a pre-load header will fail immediately.
                        Note that, even without this flag, a wrong pre-load header
                        signature check will exit in error immediately.
        -a [HASHALGO,SIGNALGO]
                        Pre-load header signature algorithm specifications.
            Current value: $SIGNATURE_SPECS_HASHALGO,$SIGNATURE_SPECS_SIGNALGO
        -n [PADDING]    Pre-load header signature padding type.
            Current value: $SIGNATURE_SPECS_PADDING
        -d [KEYPATH]    Pre-load header signature public key path.
            Current value: $SIGNATURE_SPECS_PUBKEY
        -u [KEYPATH]    Compatibility alias for '-d'.

    Other options:
        -p [PROPTYPE,PROPPATH]
                        Specify a property to retrieve from the FIT metadata.
                        PROPPATH should be the full path to the property, e.g.
                        /some/path/propname.
                        PROPTYPE should be one of:
                        - s: text string
                        - d: integer
                        - x: hexadecimal
                        - r: raw value

Examples
    $0 info u-boot.itb -p s,/version -d /my/pubkey.crt
"

usage() {
	echo "$USAGE"
}


############################
# Pre-load header stuff
############################

# $1 infile
# $2 value size in 4-bytes blocks
# $3 value offset in 4-bytes blocks
# returns value as hex string
read_header_metadata_value() {
	# Note: shellcheck is not happy about this line, but NO, using 'echo' is NOT useless: 
	# echo adds a line ending, which is critical here.
	echo $(dd if="$1" bs=4 count="$2" skip="$3" 2>/dev/null | hexdump -n 32 -e '32/1 "%02x"')
}

read_header_metadata_int() {
	echo $((0x$(read_header_metadata_value "$1" "$2" "$3")))
}

# Returns the highest 2-power divisor between $1 and $2
# Maximum value is 4k
get_highest_divisor() {
	divisor=1
	arg1=$1
	arg2=$2
	while [ $divisor -lt 4096 ]; do
		new_divisor=$((divisor * 2))
		[ $((arg1 % new_divisor)) = 0 ] || break
		[ $((arg2 % new_divisor)) = 0 ] || break
		divisor=$new_divisor
	done
	echo $divisor
}

# $1 is input file
# $2 is data offset in bytes
# $3 is data length in bytes
# Raw data blob goes to stdout
extract_raw_blob() {
	offset=$2
	length=$3
	buffer_size=$(get_highest_divisor "$offset" "$length")
	dd "if=$1" bs="$buffer_size" skip=$((offset/buffer_size)) count=$((length/buffer_size)) 2>/dev/null
}

parse_keyfile_specs() {
	keylen=$(openssl rsa -pubin -in "$SIGNATURE_SPECS_PUBKEY" -text -noout | sed -n '1,1p' | sed -E 's/.*(....) bit.*/\1/')
	export SIGNATURE_SPECS_SIGNALGO="rsa$keylen"
}

# $1 signature algorithm
# $2 signature padding mode
# $3 hash algorithm
parse_algo_properties() {
	case "$1" in
		rsa1024) signature_length=128;;
		rsa2048) signature_length=256;;
		rsa4096) signature_length=512;;
		*) echo "Unsupported signature algorithm"; exit 1;;
	esac

	case "$2" in
		pss) padding_algo_name=pss; openssl_options="$openssl_options -sigopt rsa_pss_saltlen:-2";;
		*) echo "Unsupported padding mode"; exit 1;;
	esac

	case "$3" in
		sha256) hash_algo_name=sha256; checksum_len=256;;
		sha384) hash_algo_name=sha384; checksum_len=384;;
		sha512) hash_algo_name=sha512; checksum_len=512;;
		*) echo "Unsupported hash algorithm"; exit 1;;
	esac
}

check_file_has_preload() {
	rootfs_dev_magic=$(read_header_metadata_value "$1" 1 0)
	if [ "$rootfs_dev_magic" != "$PRELOAD_MAGIC" ]; then
		return 1
	fi
}

# $1: file to check
# $2: signature specs file
# Returns the actual data offset
check_preload_header() {
	header_authent_checked=false
	data_authent_checked=false

	# Create a random ID to identify files for this session in /tmp
	session_tmpfiles_id=$(openssl rand -hex 8)
	header_metadata_sig_file="/tmp/${session_tmpfiles_id}_header_metadata.sig"
	data_sig_file="/tmp/${session_tmpfiles_id}_data.sig"

	# Parse pubkey file specs
	parse_keyfile_specs

	# Parse signature specs
	parse_algo_properties "$SIGNATURE_SPECS_SIGNALGO" "$SIGNATURE_SPECS_PADDING" "$SIGNATURE_SPECS_HASHALGO"

	# Parse header metadata and set all sizes
	full_header_size=$(read_header_metadata_int "$1" 1 2)
	header_sig_size=$signature_length
	header_sig_offset=$HEADER_METADATA_LENGHT
	data_sig_size=$signature_length
	data_sig_offset=$(read_header_metadata_int "$1" 1 4)
	data_sig_shasum=$(read_header_metadata_value "$1" $((checksum_len/8/4)) 8)
	data_size=$(read_header_metadata_int "$1" 1 3)
	data_offset=$full_header_size

	# Authenticate header metadata blob with its signature
	extract_raw_blob "$1" $header_sig_offset $header_sig_size > "$header_metadata_sig_file"
	extract_raw_blob "$1" 0 $HEADER_METADATA_LENGHT | \
		openssl dgst -verify "$SIGNATURE_SPECS_PUBKEY" \
		-$hash_algo_name -sigopt rsa_padding_mode:$padding_algo_name \
		$openssl_options -signature "$header_metadata_sig_file" > /dev/null && header_authent_checked=true
	rm "$header_metadata_sig_file"
	$header_authent_checked || return 1

	# Check data signature integrity against its checksum in metadata (always sha256)
	data_sig_actual_shasum=$(extract_raw_blob "$1" "$data_sig_offset" $data_sig_size | \
		sha256sum - | cut -d ' ' -f 1)
	[ "$data_sig_shasum" != "$data_sig_actual_shasum" ]

	# Check data against its signature in header
	extract_raw_blob "$1" "$data_sig_offset" $data_sig_size > "$data_sig_file"
	extract_raw_blob "$1" "$data_offset" "$data_size" | \
		openssl dgst -verify "$SIGNATURE_SPECS_PUBKEY" \
		-$hash_algo_name -sigopt rsa_padding_mode:$padding_algo_name \
		$openssl_options -signature "$data_sig_file" > /dev/null  && data_authent_checked=true
	rm "$data_sig_file"

	$header_authent_checked && $data_authent_checked && echo "$data_offset" || return 1
}


############################
# FIT metadata stuff
############################

# $1: input image
# $2: FIT metadata offset
# Exit in error if the provided image is not a FIT
check_file_is_fit() {
        # Retrieve FIT magic number just like a preload header value
        offset=$2
        magic=$(read_header_metadata_value "$1" 1 $((offset / 4)))
        if [ "$magic" != "d00dfeed" ]; then echo "Error: image is not a FIT" >&2; exit 1; fi
}

# $1: input image
# $2: FIT metadata offset
get_fit_totalsize() {
	# Retrieve the totalsize property of FIT, juste like a preloader header value
	# DTB header is made of 32-bit values, totalsize is 2nd.
	offset=$2
	read_header_metadata_int "$1" 1 $((offset / 4 + 1))
}

# $1 input image
# $2 FIT metadata offset
# $3 FIT property type (s, d, x, r)
# $4 FIT property path including name
get_fit_property() {
	fit_totalsize=$(get_fit_totalsize "$1" "$2")
	prop_name=$(basename "$4")
	prop_path=$(dirname "$4")
	prop_type="$3"
	[ "$prop_type" = "d" ] && prop_type=i # tweak for fdtget
	extract_raw_blob "$1" "$2" "$fit_totalsize" | \
		fdtget -t "$prop_type" - "$prop_path" "$prop_name"
}


############################
# UBI stuff
############################

ubi_hook_prerun() {
	ubiblock --create "/dev/ubi$1"
	while ! [ -e "$FILE" ] > /dev/null; do
		sleep 0;
	done
}

ubi_hook_postrun() {
	ubiblock --remove "/dev/ubi$1"
}

parsefile_ubi() {
	device=$(echo "$FILE" | sed -E 's/:.*//')
	volname=$(echo "$FILE" | sed -E 's/^ubi[0-9]+://')
	volID=$(find /sys/class/ubi/"$device" -maxdepth 2 -follow -name name -exec /bin/sh -c \
		"cat {} | grep -q '^$volname\$' && echo {} | sed -E 's,.*ubi([0-9]+_[0-9]+).*,\1,'" \;)
	if [ -z "$volID" ]; then
		echo "Error: could not find UBI volume with name $volname"
		exit 1;
	fi

	FILE="/dev/ubiblock$volID"
	if ! [ -e "$FILE" ]; then
		PRERUN_HOOKS="ubi_hook_prerun $volID"
		POSTRUN_HOOKS="ubi_hook_postrun $volID"
	fi
}

parsefile() {
	if echo "$FILE" | grep -qE "^ubi.:.*"; then
		parsefile_ubi
	fi
}


############################
# Modes
# All modes expect $FILE and $FIT_METADATA_OFFSET to be set
############################

mode_info() {
	get_fit_property "$FILE" "$FIT_METADATA_OFFSET" "$INFO_PROP_TYPE" "$INFO_PROP_PATH"
}


############################
# Main
############################

PRERUN_HOOKS=""
POSTRUN_HOOKS=""

MODE="$1"
if [ -z "$MODE" ]; then
	usage
	exit 1
fi
shift
FILE="$1"
if [ -n "$FILE" ]; then
	shift
	parsefile
fi

while getopts 'p:a:n:u:m' option "$@"; do
	case "$option" in
		a)
			SIGNATURE_SPECS_HASHALGO=$(echo "$OPTARG" | cut -d ',' -f 1)
			SIGNATURE_SPECS_SIGNALGO=$(echo "$OPTARG" | cut -d ',' -f 2)
			;;
		n) SIGNATURE_SPECS_PADDING="$OPTARG";;
		u|d) SIGNATURE_SPECS_PUBKEY="$OPTARG";;
		p)
			INFO_PROP_TYPE=$(echo "$OPTARG" | cut -d ',' -f 1)
			INFO_PROP_PATH=$(echo "$OPTARG" | cut -d ',' -f 2)
			;;
		m) PRELOAD_SIGNATURE_MANDATORY=true;;
		*) ;;
	esac
done

$PRERUN_HOOKS
if check_file_has_preload "$FILE"; then
	FIT_METADATA_OFFSET=$(check_preload_header "$FILE")
	if [ -z "$FIT_METADATA_OFFSET" ]; then exit 1; fi
else
	FIT_METADATA_OFFSET=0
	if $PRELOAD_SIGNATURE_MANDATORY; then
		echo "FIT pre-load signature header not found"
		exit 1
	fi
fi

check_file_is_fit "$FILE" "$FIT_METADATA_OFFSET"

mode_"$MODE" "$@"
result=$?
# TODO: Add a trap for post hooks to always trigger
$POSTRUN_HOOKS
exit $result
