#!/bin/bash

usage()
{
	cat << 'EOF'
usage: extract-images [<Freetz base directory>] [-be|-le] <input file>
    - Freetz base directory defaults to '.'
    - SquashFS and boot loader encoding can be '-be' (big endian) or '-le'
      (little endian). Default is to try both types.
    - Input file (e.g. recover exe or mtdblock dump) to be searched for
      kernel.image and/or urlader.image
EOF
}

# Print usage info to stdout (no error, just info)
if [ $# -eq 0 ]; then
	usage
	exit 0
fi

# Parameterized helper function for 'extract_bootloader' doing the actual work
do_extract_bl()
{
	local ENDIAN_TYPE="$1"

	# Find bootloader candidates
	declare -a bootloader_offs=($(getDecOffsetsOfAllMatches "$input_file" bin "${BOOTLOADER_MAGIC[${ENDIAN_TYPE}]}"))

	# Find bootloader candidates and isolate their offsets
	for (( i=0; i<${#bootloader_offs[@]}; i++ )); do
		# For each bootloader candidate create a new numbered 'urlader.image' file
		output_file="$output_dir/urlader.image${count}"
		offs=${bootloader_offs[i]}
		# Bootloader files are either 64K or 128K long. Try both variants.
		tail -c +$(($offs+1)) "$input_file" | head -c +65536 > "$output_file.64k"
		tail -c +$(($offs+1)) "$input_file" | head -c +131072 > "$output_file.128k"
		# Identify the right candidate by grepping for known bootloader text strings
		local candidate_found=0
		for candidate in "$output_file.64k" "$output_file.128k"; do
			if [ $candidate_found -eq 1 ]; then
				rm -f "$candidate"
				continue
			fi
			if grep -iq 'entering passive mode' "$candidate" && grep -iq 'ftp.\? server ready' "$candidate"; then
				mv "$candidate" "$output_file"
				echo "    $output_file - loader @ $offs, $ENDIAN_TYPE endian"
				n=$((n+1))
				count="-$n"
				candidate_found=1
			else
				# If candidate does not match, remove it
				rm -f "$candidate"
			fi
		done
	done
}

# Extracts urlader.image candidates from input file, based on bootloader magic
# string. Applies heuristic check for typical strings occurring in AVM
# bootloaders. Detects older ADAM as well as newer EVA bootloaders.
extract_bootloader()
{
	echo -e "\nExtracting bootloader (urlader.image) ..."
	n=1
	unset count
	[ "$endian" == "-any" -o "$endian" == "-le" ] && \
		do_extract_bl "little"
	[ "$endian" == "-any" -o "$endian" == "-be" ] && \
		do_extract_bl "big"
	[ $n -eq 1 ] && echo "    No bootloader candidate found in input file"
}

get_ext2fs_length()
{
	local ext2fs_offset="$1"
	local ENDIAN_TYPE="$2"   # unused

	local rle_stream_length

	rle_stream_length=$(tail -c +$(($ext2fs_offset + 5)) "$input_file" | $mod_tools/avm-rle-stream-length)
	[ $? -eq 0 ] && echo "$(($rle_stream_length + 4))" || echo 0 # dummy length
}

get_squashfs_length()
{
	local squashfs_offset="$1"
	local ENDIAN_TYPE="$2"

	local squashfs_length=, nmi_offset=, nmi_length=

	local NMI_BOOT_PATTERN="4E4D4920426F6F7420566563746F7200"

	# Note: 4 bytes of length information are found at (SquashFS magic offset + 8)
	#       valid for all vanilla/AVM SquashFS versions up to (including) SquashFS-3.x
	#       NOT valid for vanilla SquashFS-4.x but still valid for AVM-SquashFS-4.x
	#       because of AVM's "mkfs_time=bytes_used" modification (s. SquashFS packages for details)
	squashfs_length=$(getHexContentAtOffset "$input_file" $((${squashfs_offset} + 8)) 4 | { [ "$ENDIAN_TYPE" == "little" ] && invertEndianness || cat; })
	squashfs_length=$((0x${squashfs_length}))

	declare -a nmi_offsets=($(getDecOffsetsOfAllMatches "$input_file" bin "$NMI_BOOT_PATTERN"))
	for (( i=0; i<${#nmi_offsets[@]}; i++ )); do
		nmi_offset="${nmi_offsets[i]}"
		if [ "${squashfs_offset}" -le "${nmi_offset}" -a "${nmi_offset}" -lt "$((${squashfs_offset} + ${squashfs_length}))" ]; then
			# NMI pattern is found within SquashFS

			# NMI vector length is encoded within the NMI block 4 bytes before 'NMI Boot'
			nmi_length="$(getHexContentAtOffset "$input_file" $((${nmi_offset} - 4)) 4)"
			nmi_length=$((0x${nmi_length}))
			if [ "${nmi_length}" -eq 256 -o "${nmi_length}" -eq 4096 ]; then
				# NMI_length is one of the known correct values => add NMI vector length to SquashFS length
				squashfs_length=$((${squashfs_length} + ${nmi_length}))
				break;
			fi
		fi
	done
	echo "${squashfs_length}"
}

# Parameterized helper function for 'extract_kernel_filesystem' doing the actual work
do_extract_kfs()
{
	local ENDIAN_TYPE="$1"

	local rootfs_type

	declare -A rootfs_magic=(
		[ext2fs]="${EXT2FS_MAGIC[${ENDIAN_TYPE}]}"
		[squashfs]="${SQUASHFS_MAGIC[${ENDIAN_TYPE}]}"
	)
	for rootfs_type in ext2fs squashfs; do
		declare -a rootfs_offs=($(getDecOffsetsOfAllMatches "$input_file" bin "${rootfs_magic[${rootfs_type}]}"))
		if [ ${#rootfs_offs[@]} -gt 0 ]; then
			break;
		fi
	done

	for (( i=0; i<${#rootfs_offs[@]}; i++ )); do
		# Isolate SquashFS length candidates (considering endian type)
		rootfs_length[i]=$(get_${rootfs_type}_length "${rootfs_offs[i]}" "$ENDIAN_TYPE")

		# Special case: filesystem length 0 can happen if a raw (mtdblock) dump
		# from the box is processed, because the box seems to erase the length
		# info when flashing.
		[ ${rootfs_length[i]} -eq 0 ] && rootfs_length[i]=$((input_file_size-rootfs_offs[i]))
	done

	# Find kernel candidates
	declare -a kernel_offs=($(getDecOffsetsOfAllMatches "$input_file" bin "$KERNEL_MAGIC"))

	for (( i=0; i<${#kernel_offs[@]}; i++ )); do
		# Isolate kernel length candidates (always little endian)
		# Note: 4 bytes of length information are found at (kernel magic offset + 4)
		tmp=$(getHexContentAtOffset "$input_file" $((${kernel_offs[i]} + 4)) 4 | invertEndianness)
		kernel_length[i]=$(printf "%d" "0x$tmp")
		# Pad kernel length to 256
		kernel_pad=$(( (256 - kernel_length[i] % 256) % 256 ))
		kernel_length[i]=$(( kernel_length[i] + kernel_pad ))
	done

	# Check all combinations of kernel vs SquashFS offset candidates to find plausible ones
	for (( k=0; k<${#kernel_offs[@]}; k++ )); do
		# Sanity check #1: kernel offset + length must not exceed total file size
		[ $((kernel_offs[k] + kernel_length[k])) -le $input_file_size ] || continue
		for (( s=0; s<${#rootfs_offs[@]}; s++ )); do
			# Sanity check #2: SquashFS offset + length must not exceed total file size
			[ $((rootfs_offs[s] + rootfs_length[s])) -le $input_file_size ] || continue

			# Sanity check #3: kernel + SquashFS lengths must be 85-100% of total file size
			percentage=$(( (kernel_length[k] + rootfs_length[s]) * 100 / input_file_size ))
			[ $percentage -ge 85 -a $percentage -le 100 ] || continue

			# Check layout: does SquashFS immediately follow kernel?
			gap=$((rootfs_offs[s] - kernel_length[k] - kernel_offs[k]))

			# workaround for some images with incorrect(?) or at least strange kernel_length
			# known affected images: FRITZ.Box_Fon_WLAN_7390.AnnexB.06.03.recover-image.exe
			if [ $gap -eq 256 ]; then
				# treat gap of 256 bytes as zero-gap and increase kernel_length by 256 bytes,
				# i.e. pad kernel with additional 256 bytes
				kernel_length[k]=$((kernel_length[k] + 256))
				gap=0
			fi

			if [ $gap -eq 0 ]; then
				# Hidden root (SquashFS contained in kernel.image)
				file_length=$((kernel_length[k] + rootfs_length[s]))
				output_file="$output_dir/hr_kernel.image${count}"
				echo "    $output_file - kernel @ ${kernel_offs[k]}, $ENDIAN_TYPE endian SquashFS @ ${rootfs_offs[s]}, total length = $file_length"
				echo "        This is a \"hidden root\" image (compound kernel + SquashFS)."
				echo "        For your convenience I will split them for you so you do not have to use tools/find-squashfs"
				tail -c +$((kernel_offs[k]+1)) "$input_file" | head -c +$file_length > "$output_file"
			fi

			# Sanity check #4: kernel and SquashFS must not intersect
			[ $gap -lt 0 ] && gap=$((kernel_offs[k] - rootfs_length[s] - rootfs_offs[s]))
			[ $gap -lt 0 ] && continue

			# Plausible candidates for kernel and filesystem
			# Save kernel and filesystem as separate files

			# kernel.image
			output_file="$output_dir/kernel.image${count}"
			echo "    $output_file - kernel @ ${kernel_offs[k]}, length = ${kernel_length[k]}"
			echo "        If you are interested in an uncompressed kernel, use tools/unpack-kernel"
			tail -c +$((kernel_offs[k]+1)) "$input_file" | head -c +${kernel_length[k]} > "$output_file"

			# filesystem.image
			output_file="$output_dir/filesystem.image${count}"
			echo "    $output_file - $ENDIAN_TYPE endian ${rootfs_type} @ ${rootfs_offs[s]}, length = ${rootfs_length[s]}"
			echo "        If you want to see the unpacked filesystem, try tools/build/bin/fakeroot '--' tools/unsquashfs*"
			echo "        The unpacked SquashFS may contain filesystem_core.squashfs which you can then also unsquash."
			if [ "${rootfs_type}" == "ext2fs" ]; then
				tail -c +$((rootfs_offs[s]+1)) "$input_file" | head -c 4 > "$output_file"
				tail -c +$((rootfs_offs[s]+5)) "$input_file" | "$mod_tools/avm-rle-decode" >> "$output_file"
				tail -c +257 "$output_file" > "$output_file.ext2"
			else
				tail -c +$((rootfs_offs[s]+1)) "$input_file" | head -c +${rootfs_length[s]} > "$output_file"
			fi

			n=$((n+1))
			count="-$n"
		done
	done
}

# Extracts kernel.image (and possibly filesystem.image) candidates from input
# file, based on kernel and SquashFS magic strings. Applies heuristic sanity
# checks for offsets, sizes and gaps between kernel and filesystem.
extract_kernel_filesystem()
{
	echo -e "\nExtracting kernel + filesystem ..."
	n=1
	unset count
	[ "$endian" == "-any" -o "$endian" == "-le" ] && \
		do_extract_kfs "little"
	[ "$endian" == "-any" -o "$endian" == "-be" ] && \
		do_extract_kfs "big"
	[ $n -eq 1 ] && echo "    No kernel + filesystem candidates found in $input_file"
}

# Parameter handling
mod_tools="$(dirname $(readlink -f ${0}))"
endian="-any"
case $# in
	0)
		usage
		exit 0
		;;
	1)
		input_file="$1"
		;;
	2)
		[ "$1" == "-be" -o "$1" == "-le" ] && endian="$1" || mod_tools="${1%/}/tools"
		input_file="$2"
		;;
	3)
		mod_tools="${1%/}/tools"
		[ "$2" == "-be" -o "$2" == "-le" ] && endian="$2" || { usage >&2; exit 1; }
		input_file="$3"
		;;
	*)
		usage >&2
		exit 1
		;;
esac

source "$mod_tools/freetz_bin_functions"

# Magic byte sequences (hex) to identify kernel, SquashFS and bootloader
KERNEL_MAGIC="8112EDFE"
#
SQUASHFS_MAGIC_BE="73717368"
SQUASHFS_MAGIC_LE=$(echo "$SQUASHFS_MAGIC_BE" | invertEndianness)
declare -A SQUASHFS_MAGIC=(
	[big]=${SQUASHFS_MAGIC_BE}
	[little]=${SQUASHFS_MAGIC_LE}
)
#
# sqsh followed by AVM-RLE-encoded sequence of (252 + 1024) zero bytes, 252: AVM-zero-header, 1024: first 1024 EXT2FS bytes
# opcode: 0x81, payloadLen: 0xFC04 (little-endian), payload: 0x00
declare -A EXT2FS_MAGIC=(
	[big]="${SQUASHFS_MAGIC_BE}81FC0400"
	[little]="${SQUASHFS_MAGIC_LE}81FC0400"
)
#
BOOTLOADER_MAGIC_BE="40809000"
declare -A BOOTLOADER_MAGIC=(
	[big]=${BOOTLOADER_MAGIC_BE}
	[little]=$(echo "$BOOTLOADER_MAGIC_BE" | invertEndianness)
)

input_file_size=$(stat -L -c %s "$input_file")

output_dir="$(basename "$input_file").unp"
rm -rf "$output_dir"
mkdir -p "$output_dir"

extract_bootloader
extract_kernel_filesystem