#!/bin/bash

#
# Remove NMI vector from AVM image files
# Author: Alexander Kriegisch (http://scrum-master.de) and Freetz developer team
#
# NMI vector v1:
#  1. starts with byte sequence 3C1A8000375A038003400008
#  2. located at (fixed) offset 0xBE0000
#  3. has (fixed) length of 256 bytes
#  4. NMI vector size is saved in 4 bytes at (fixed) offset 0xBE0044,
#     the same offset could also be determined by locating the string "NMI Boot"
#     and subtracting 4 bytes from its offset
#  5. image size without the NMI vector and without the image checksum is saved
#     in 4 bytes at (fixed) offset 0xBE0040,
#     the same offset could also be determined by locating the string "NMI Boot"
#     and subtracting 8 bytes from its offset
#
# NMI vector v2:
#  1. starts with byte sequence 409AE8053C1A8000375A038003400008 (v1 sequence prefixed with 409AE805)
#  2. same as v1
#  3. has (fixed) length of 4096 bytes
#  4. same as v1
#  5. same as v1
#
# NMI vector v3:
#  1. starts with byte sequence 409AE805409BE0043C1A8000375A0380
#  2. same as v1
#  3. same as v2
#  4. same as v1
#  5. same as v1
#

source "$(dirname $(readlink -f ${0}))/freetz_bin_functions"

INPUT_FILE="$1"
OUTPUT_FILE="$2"

if [ $# -ne 2 -o ! -e "$INPUT_FILE" ]; then
	echo "Usage: ${0##*/} input.image output.image"
	exit 2
fi

# see linux-${KERNEL_VERSION}/arch/mips/fusiv/fusiv_mips32/fusiv_mtd.c (7390) for details
NMI_BOOT_STRING="NMI Boot Vector"
NMI_BOOT_PATTERN="4E4D4920426F6F7420566563746F7200" # the string above (zero-terminated) as sfk-bin-pattern
expectedNmiBootOffset=0xBE0048

nmiVectorOffset=0xBE0000
declare -A nmiVectorKnownStartSequence=(
		[v1]=3C1A8000375A038003400008
	[v2]=409AE8053C1A8000375A038003400008
	[v3]=409AE805409BE0043C1A8000375A0380
)

getInputFileSize() {
	stat -L -c "%s" "$INPUT_FILE"
}


declare -a nmiBootOffsets=($(getHexOffsetsOfAllMatches "$INPUT_FILE" bin "${NMI_BOOT_PATTERN}"))

if [ ${#nmiBootOffsets[@]} -eq 0 ]; then
	echo "No NMI vector found" >&2
	exit 1
fi

if [ ${#nmiBootOffsets[@]} -gt 1 ]; then
	echo "[Error] Multiple '${NMI_BOOT_STRING}'-matches found: ${nmiBootOffsets[@]}" >&2
	exit 2
fi

nmiBootOffset="${nmiBootOffsets[0]}"
if [ "${nmiBootOffset}" != "${expectedNmiBootOffset}" ]; then
	echo "[Error] '${NMI_BOOT_STRING}'-match found at unexpected offset ${nmiBootOffset}" >&2
	exit 2
fi

# NMI vector length is encoded within the NMI block 4 bytes before 'NMI Boot'
nmiVectorLength="0x$(getHexContentAtOffset "$INPUT_FILE" $((${nmiBootOffset} - 4)) 4)"
if [ $((${nmiVectorLength})) -ne 256 -a $((${nmiVectorLength})) -ne 4096 ]; then
	echo "[Error] Unexpected/unsupported NMI vector length: $((${nmiVectorLength}))" >&2
	exit 2
fi

# check if the image contains the TI-checksum at the end
TI_CHECKSUM_MAGIC_BE="23DE53C4"
if [ "$(getHexContentAtOffset "$INPUT_FILE" $(($(getInputFileSize) - 8)) 4)" == "${TI_CHECKSUM_MAGIC_BE}" ]; then
	# TI checksum found
	imageChecksumSize=8
else
	# no TI checksum found
	imageChecksumSize=0
fi

# real image size is encoded within the NMI block 8 bytes before 'NMI Boot'
imageSizeWithinNMI="0x$(getHexContentAtOffset "$INPUT_FILE" $((${nmiBootOffset} - 8)) 4)"

if [ "$((${imageSizeWithinNMI}))" -gt "$((${nmiVectorOffset}))" ]; then
	# NMI boot block resides within the (filesystem) image, typically Fritz!Box
	# kernel.image | filesystem.image part1 | NMI block | filesystem.image part2
	imageSizeWithinNMIPadded256=$(( ${imageSizeWithinNMI} + (256 - ${imageSizeWithinNMI} % 256) % 256 ))
	imageSizeExpected=$(printf "0x%08X" $(( $(getInputFileSize) - ${nmiVectorLength} - ${imageChecksumSize} )))
	if [ "$((${imageSizeWithinNMI}))" -ne "$((${imageSizeExpected}))" -a "$((${imageSizeWithinNMIPadded256}))" -ne "$((${imageSizeExpected}))" ]; then
		echo "[Error] Image size encoded in NMI block ($((${imageSizeWithinNMI}))) doesn't match calculated size ($((${imageSizeExpected})))" >&2
		exit 2
	fi
else
	# NMI boot block resides after the (filesystem) image, typically Fritz!WLAN Repeater / Fritz!Powerline Adapter
	# kernel.image | filesystem.image | gap padding bytes | NMI block
	fileSize=$(getInputFileSize)
	fileSizeExpected=$(printf "0x%08X" $(( ${nmiVectorOffset} + ${nmiVectorLength} + ${imageChecksumSize} )))
	if [ "$((${fileSize}))" -ne "$((${fileSizeExpected}))" ]; then
		echo "[Error] File size ($((${fileSize}))) doesn't match the expected one ($((${fileSizeExpected})))" >&2
		exit 2
	fi
fi

# get first 32 bytes of the NMI vector
nmiVectorStartSequence=$(getHexContentAtOffset "$INPUT_FILE" ${nmiVectorOffset} 32)
# and check if they contain one of the known sequences
for v in v1 v2 v3; do
	if [ "${nmiVectorStartSequence:0:${#nmiVectorKnownStartSequence[${v}]}}" != "${nmiVectorKnownStartSequence[${v}]}" ]; then
		continue
	fi

	echo -n "NMI vector ${v} found at offset ${nmiVectorOffset}, removing it ... "
	{
		head -c "$((${nmiVectorOffset}))" "$INPUT_FILE"
		tail -c "+$((${nmiVectorOffset} + ${nmiVectorLength} + 1))" "$INPUT_FILE"
	} | head -c "$((${imageSizeWithinNMI}))" > "$OUTPUT_FILE"
	echo "done."

	# debug: dump NMI vector to a separate file
	#tail -c "+$((${nmiVectorOffset} + 1))" "$INPUT_FILE" | head -c $((${nmiVectorLength})) > "$INPUT_FILE".nmi-vector

	exit 0
done

echo "[Error] NMI vector start sequence (${nmiVectorStartSequence}) doesn't match any of the known sequences" >&2
exit 2