#
# Copyright (C) 2025 Amlogic, Inc. All rights reserved.
#

#!/bin/bash

EXEC_BASEDIR=$(dirname $(readlink -f $0))
BASEDIR_TOP=$(readlink -f ${EXEC_BASEDIR}/..)
ACPU_IMAGETOOL=${EXEC_BASEDIR}/../binary-tool/acpu-imagetool
AMFCTOOL=${EXEC_BASEDIR}/../../tools/aml_zstd

#
# Settings
#
#BASEDIR_BUILD="${BASEDIR_TOP}/output"
BASEDIR_BUILD="${BASEDIR_TOP}/`date +%Y%m%d%H%M%S%N`"
postfix=.signed

declare -a BLX_BIN_MIN_SIZE=("65536" "32768")
declare -a BLX_BIN_MAX_SIZE=("524288" "1048576")

input_bl31=""
input_bl32=""
input_bl3_tee=""

BL31_SIZE=0
BL32_SIZE=0

function align_4k() {
	local file=$1
	local raw_size=
	local file_size=
	local rem_val=

	raw_size=`stat -c %s ${file}`
	file_size=`expr ${raw_size} + 4095`
	rem_val=`expr ${file_size} % 4096`
	file_size=`expr ${file_size} - ${rem_val}`
	#echo ==== align file ${file} from ${raw_size} to ${file_size} ====
	dd if=/dev/zero of=${file}_tmp.bin bs=1 count=${file_size} &> /dev/null
	dd if=${file} of=${file}_tmp.bin bs=1 count=${file_size} conv=notrunc &> /dev/null
	mv ${file}_tmp.bin ${file}
}

function compress_binary() {
	local path=$1
	local comp_lv=$2
	local blx=$3
	amfc_zstd_hdr=$1/amfc_zstd_hdr.bin
	${AMFCTOOL} $1/$3-payload.bin -$comp_lv -o $1/$3-payload.bin.zstd
	bin_org_size=`stat -c %s $1/$3-payload.bin`
	bin_zstd_size=`stat -c %s $1/$3-payload.bin.zstd`
	printf "%s" "@ZST" >  $amfc_zstd_hdr

	printf "%02x%02x%02x%02x" $[(bin_org_size) & 0xff] \
	$[((bin_org_size) >> 8) & 0xff] $[((bin_org_size) >> 16) & 0xff] \
	$[((bin_org_size) >> 24) & 0xff] | xxd -r -ps >>  $amfc_zstd_hdr

	printf "%02x%02x%02x%02x" $[(bin_zstd_size) & 0xff] \
	$[((bin_zstd_size) >> 8) & 0xff] $[((bin_zstd_size) >> 16) & 0xff] \
	$[((bin_zstd_size) >> 24) & 0xff] | xxd -r -ps >>  $amfc_zstd_hdr

	cat $amfc_zstd_hdr $1/$3-payload.bin.zstd > $1/$3-payload.bin.tmp

	bin_zstd_size=`stat -c %s $1/$3-payload.bin.tmp`

	if [ "$blx" == "bl31" ]; then
		if [ "${bin_zstd_size}" -lt "${BLX_BIN_MIN_SIZE[0]}" ]; then
			dd if=/dev/zero of=${BASEDIR_BUILD}/$3-payload.bin bs=${BLX_BIN_MIN_SIZE[0]} count=1  &> /dev/null
			dd if=$1/$3-payload.bin.tmp of=$1/$3-payload.bin conv=notrunc  &> /dev/null
		elif [ "${bin_zstd_size}" -gt "${BLX_BIN_MAX_SIZE[0]}" ]; then
			echo "bl31 size exceed 512k"
			exit 1
		else
			align_4k $1/$3-payload.bin.tmp
			cp $1/$3-payload.bin.tmp $1/$3-payload.bin
		fi
	elif [ "$blx" == "bl32" ]; then
		if [ "${bin_zstd_size}" -lt "${BLX_BIN_MIN_SIZE[1]}" ] || [ "fastboot" == "${boot_mode}" ]; then
			dd if=/dev/zero of=${BASEDIR_BUILD}/$3-payload.bin bs=${BLX_BIN_MIN_SIZE[1]} count=1  &> /dev/null
			dd if=$1/$3-payload.bin.tmp of=$1/$3-payload.bin conv=notrunc  &> /dev/null
		elif [ "${bin_zstd_size}" -gt "${BLX_BIN_MAX_SIZE[1]}" ]; then
			echo "bl32 size exceed 1024k"
			exit 1
		else
			align_4k $1/$3-payload.bin.tmp
			cp $1/$3-payload.bin.tmp $1/$3-payload.bin
		fi
	fi
	rm $amfc_zstd_hdr -f
}

function check_file() {
	if [ ! -f "$1" ]; then
		echo "Error: file \""$1"\" does NOT exist"
		exit 1
	fi
}

function sign_tee() {
	local argv=("$@")
	local i=0

	# Parse args

	while [ $i -lt $# ]; do
		arg="${argv[$i]}"
		#echo "i=$i argv[$i]=${argv[$i]}"
		i=$((i + 1))
		case "$arg" in
			--bl31)
				input_bl31="${argv[$i]}" ;;
			--bl32)
				input_bl32="${argv[$i]}" ;;
			--bl3_tee)
				input_bl3_tee="${argv[$i]}" ;;
			--output)
				output="${argv[$i]}" ;;
			--chipset_name)
				chipset_name="${argv[$i]}" ;;
			--chipset_variant)
				chipset_variant="${argv[$i]}" ;;
			--key_type)
				key_type="${argv[$i]}" ;;
			--soc)
				soc="${argv[$i]}" ;;
			--build_type)
				build_type="${argv[$i]}" ;;
			--extra_args)
				boot_mode="${argv[$i]}" ;;
			*)
				echo "Unknown option $arg"; exit 1
				;;
		esac
		i=$((i + 1))
	done

	# At least one of bl31 or bl32 must exist
	if [ -z ${input_bl31} ] && [ -z ${input_bl32} ]; then
		echo "input missing 31 or 32 ${input_bl31} ${input_bl32}"
		exit 1
	fi
	if [ ! -z ${input_bl31} ]; then
		check_file "$input_bl31"
	fi
	if [ ! -z ${input_bl32} ]; then
		check_file "$input_bl32"
	fi
	if [ -z ${input_bl3_tee} ]; then
		if [ -z ${input_bl31} ] || [ -z ${input_bl32} ]; then
			# If bl3_tee does not exist, both bl31 and bl32 must exist.
			echo "input ${input_bl3_tee} invalid"
			exit 1
		fi
	fi

	if [ -z ${output} ]; then
		echo "output ${output} invalid"
		exit 1
	fi

	if [ -z ${chipset_name} ]; then
		chipset_name="n611x1"
	fi

	if [ -z ${key_type} ]; then
		key_type="dev-keys"
	fi

	if [ -z ${soc} ]; then
		soc="c5"
	fi

	mkdir ${BASEDIR_BUILD}

	echo "BASEDIR_BUILD is ${BASEDIR_BUILD}"

	if [ "fastboot" == "${boot_mode}" ]; then
		BL32_SIZE=32768
	fi

	if [ -f "${input_bl31}" ]; then
		cp ${input_bl31} ${BASEDIR_BUILD}/bl31-payload.bin -f
		compress_binary ${BASEDIR_BUILD} 9 bl31
		BL31_SIZE=`stat -c %s ${BASEDIR_BUILD}/bl31-payload.bin`
	fi

	if [ -f "${input_bl32}" ]; then
		cp ${input_bl32} ${BASEDIR_BUILD}/bl32-payload.bin -f
		compress_binary ${BASEDIR_BUILD} 9 bl32
		BL32_SIZE=`stat -c %s ${BASEDIR_BUILD}/bl32-payload.bin`
	fi

	if [ -f "${input_bl3_tee}" ]; then
		cp ${input_bl3_tee} ${BASEDIR_BUILD}/blob-bl3-tee.bin
	fi

	${EXEC_BASEDIR}/download-keys.sh ${key_type} ${soc} chipset

	#
	# Settings
	#

	BASEDIR_AESKEY_PROT_BL2="${BASEDIR_TOP}/keys/${key_type}/${soc}/chipset/bl2/aes/${chipset_name}${EXT_SUFFIX}"
	BASEDIR_RSAKEY_LVLX_BL2="${BASEDIR_TOP}/keys/${key_type}/${soc}/chipset/bl2/rsa/${chipset_name}${EXT_SUFFIX}"

	BASEDIR_AESKEY_PROT_BL31="${BASEDIR_TOP}/keys/${key_type}/${soc}/chipset/bl31/aes/${chipset_name}${EXT_SUFFIX}"
	BASEDIR_RSAKEY_LVLX_BL31="${BASEDIR_TOP}/keys/${key_type}/${soc}/chipset/bl31/rsa/${chipset_name}${EXT_SUFFIX}"

	BASEDIR_AESKEY_PROT_BL32="${BASEDIR_TOP}/keys/${key_type}/${soc}/chipset/bl32/aes/${chipset_name}${EXT_SUFFIX}"
	BASEDIR_RSAKEY_LVLX_BL32="${BASEDIR_TOP}/keys/${key_type}/${soc}/chipset/bl32/rsa/${chipset_name}${EXT_SUFFIX}"

	BASEDIR_OUTPUT_BLOB=${output}
	postfix=.signed

	#
	# Arguments
	#

	_BASEDIR_AESKEY_PROT_DIR_BL31="BASEDIR_AESKEY_PROT_BL31"
	BASEDIR_AESKEY_PROT_DIR_BL31=${!_BASEDIR_AESKEY_PROT_DIR_BL31}

	_BASEDIR_AESKEY_PROT_DIR_BL32="BASEDIR_AESKEY_PROT_BL32"
	BASEDIR_AESKEY_PROT_DIR_BL32=${!_BASEDIR_AESKEY_PROT_DIR_BL32}

	EXEC_ARGS="${EXEC_ARGS}"

	### Input: payload ###
	if [ -f "${input_bl31}" ]; then
		EXEC_ARGS="${EXEC_ARGS} --infile-bl31-payload=${BASEDIR_BUILD}/bl31-payload.bin"
		EXEC_ARGS="${EXEC_ARGS} --infile-aes256-bl31-payload=${BASEDIR_AESKEY_PROT_DIR_BL31}/genkey-prot-bl31.bin"
		EXEC_ARGS="${EXEC_ARGS} --size-payload-bl31=$BL31_SIZE"
	fi

	if [ -f "${input_bl32}" ]; then
		EXEC_ARGS="${EXEC_ARGS} --infile-bl32-payload=${BASEDIR_BUILD}/bl32-payload.bin"
		EXEC_ARGS="${EXEC_ARGS} --infile-aes256-bl32-payload=${BASEDIR_AESKEY_PROT_DIR_BL32}/genkey-prot-bl32.bin"
		EXEC_ARGS="${EXEC_ARGS} --size-payload-bl32=$BL32_SIZE"
	fi

	if [ -f "${input_bl3_tee}" ]; then
		EXEC_ARGS="${EXEC_ARGS} --infile-blob-bl3-tee=${BASEDIR_BUILD}/blob-bl3-tee.bin"
	fi

	### Input: Chipset Level-1/2 Private RSA keys
	EXEC_ARGS="${EXEC_ARGS} --infile-signkey-chipset-lvl2=${BASEDIR_RSAKEY_LVLX_BL2}/level-2-rsa-priv.pem"

	### Output: blobs ###
	EXEC_ARGS="${EXEC_ARGS} --outfile-blob-bl3-tee=${BASEDIR_OUTPUT_BLOB}"

	### boot blobs args
	EXEC_ARGS="${EXEC_ARGS} --scs-family=${soc}"

	echo "$EXEC_ARGS"

	${ACPU_IMAGETOOL} \
		create-boot-blobs \
		${EXEC_ARGS}

	#cp ${BASEDIR_BUILD}/blob-bl3-tee.bin${postfix} $output
}

rm -rf ${BASEDIR_BUILD}
sign_tee $@
rm -rf ${BASEDIR_BUILD}