#!/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

# Parametrised helper function for 'extract_bootloader' doing the actual work
do_extract_bl()
{
	local BOOTLOADER_MAGIC="$1"
	local ENDIAN_TYPE="$2"
	# Find bootloader candidates and isolate their offsets
	for offs in $($mod_base/tools/sfk hexfind "$input_file" -bin /$BOOTLOADER_MAGIC/ | sed -nr 's/.*hit at offset (.*)/\1/p' | xargs printf "%d\n"); do
		# For each bootloader candidate create a new numbered 'urlader.image' file
		output_file="$output_dir/urlader.image-$n"
		# Bootloader files are always exactly 64K long
		tail -c +$(($offs+1)) "$input_file" | head -c +65536 > "$output_file"
		# Identify the right candidate by grepping for known bootloader text strings
		if grep -iq 'entering passive mode' "$output_file" && grep -iq 'ftp.\? server ready' "$output_file"; then
			echo "    $output_file - loader @ $offs, $ENDIAN_TYPE endian"
			n=$((n+1))
		else
			# If candidate does not match, remove it
			rm -f "$output_file"
		fi
	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
	[ "$endian" == "-any" -o "$endian" == "-le" ] && \
		do_extract_bl "$BOOTLOADER_MAGIC_LE" "little"
	[ "$endian" == "-any" -o "$endian" == "-be" ] && \
		do_extract_bl "$BOOTLOADER_MAGIC_BE" "big"
	[ $n -eq 1 ] && echo "    No bootloader candidate found in input file"
}


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

	# Find SquashFS candidates ('sfk hexfind' dumps offsets and hexdump contexts)
	hexfind_dump="$($mod_base/tools/sfk hexfind "$input_file" -bin /$SQUASHFS_MAGIC/)"
	# Isolate arrays of SquashFS offset candidates (hex & dec)
	squashfs_offs_hex=($(echo "$hexfind_dump" | sed -nr 's/.*hit at offset 0x(.*)/\1/p'))
	squashfs_offs=($(echo "$hexfind_dump" | sed -nr 's/.*hit at offset (.*)/\1/p' | xargs printf "%d\n"))
	for (( i=0; i<${#squashfs_offs[@]}; i++ )); do
		# Isolate array of SquashFS length candidates from hexdump lines, considering endian type
		# Note: 4 bytes of length information are found at (SquashFS magic offset + 8)
		tmp=$(echo "$hexfind_dump" | grep " 0*${squashfs_offs_hex[i]}\$" | sed -r 's/([^ ]+ ){2}([^ ]+).*/\2/')
		[ "$ENDIAN_TYPE" == "big" ] \
			&& squashfs_length[i]=$(printf "%d" "0x$tmp") \
			|| squashfs_length[i]=$(printf "%d" "0x${tmp:6:2}${tmp:4:2}${tmp:2:2}${tmp:0:2}")
	done

	# Find kernel candidates ('sfk hexfind' dumps offsets and hexdump contexts)
	hexfind_dump="$($mod_base/tools/sfk hexfind "$input_file" -bin /$KERNEL_MAGIC/)"
	# Isolate arrays of kernel offset candidates (hex & dec)
	kernel_offs_hex=($(echo "$hexfind_dump" | sed -nr 's/.*hit at offset 0x(.*)/\1/p'))
	kernel_offs=($(echo "$hexfind_dump" | sed -nr 's/.*hit at offset (.*)/\1/p' | xargs printf "%d\n"))
	for (( i=0; i<${#kernel_offs[@]}; i++ )); do
		# Isolate array of kernel length candidates from hexdump lines (always little endian)
		# Note: 4 bytes of length information are found at (kernel magic offset + 4)
		tmp=$(echo "$hexfind_dump" | grep " 0*${kernel_offs_hex[i]}\$" | sed -r 's/([^ ]+ )([^ ]+).*/\2/')
		kernel_length[i]=$(printf "%d" "0x${tmp:6:2}${tmp:4:2}${tmp:2:2}${tmp:0:2}")
	done

	# Check all combinations of kernel vs SquashFS offset candidates to find plausible ones
	for (( k=0; k<${#kernel_offs[@]}; k++ )); do
		for (( s=0; s<${#squashfs_offs[@]}; s++ )); do
			# Plausible gap between kernel and subsequent SquashFS is in [0..256] (kernel padding)
			kernel_pad=$((squashfs_offs[s] - kernel_length[k] - kernel_offs[k]))
			[ $kernel_pad -le 0 -o $kernel_pad -gt 256 ] && continue
			# Determine image file length
			file_length=$((kernel_length[k] + kernel_pad + squashfs_length[s]))
			output_file="$output_dir/kernel.image-$((n++))"
			# Create a new numbered 'kernel.image' file
			echo "    $output_file - kernel @ ${kernel_offs[k]}, $ENDIAN_TYPE endian SquashFS @ ${squashfs_offs[s]}, total length = $file_length"
			tail -c +$((kernel_offs[k]+1)) "$input_file" | head -c +$file_length > "$output_file"
			# Add checksum (another 8 bytes at the end of the file)
			"$mod_base/tools/tichksum" "$output_file" | sed 's/^/        /'
			break
		done
	done
}

# Extracts kernel.image candidates from input file, based on kernel and
# SquashFS magic strings. Applies heuristic check if kernel end offset
# calculated from assumed length header field is close enough to subsequent
# SquashFS start offset. Detects hidden root images used in Freetz-enabled
# firmwares, but currently no other image types.
extract_kernel_filesystem()
{
	echo -e "\nExtracting hidden root kernel + filesystem (kernel.image) ..."
	n=1
	[ "$endian" == "-any" -o "$endian" == "-le" ] && \
		do_extract_kfs "$SQUASHFS_MAGIC_LE" "little"
	[ "$endian" == "-any" -o "$endian" == "-be" ] && \
		do_extract_kfs "$SQUASHFS_MAGIC_BE" "big"
	[ $n -eq 1 ] && echo "    No hidden root kernel + filesystem candidate found in input file"
}

# Parameter handling
mod_base="."
endian="-any"
case $# in
	0)
		usage
		exit 0
		;;
	1)
		input_file="$1"
		;;
	2)
		[ "$1" == "-be" -o "$1" == "-le" ] && endian="$1" || mod_base="$1"
		input_file="$2"
		;;
	3)
		mod_base="$1"
		[ "$2" == "-be" -o "$2" == "-le" ] && endian="$2" || { usage >&2; exit 1; }
		input_file="$3"
		;;
	*)
		usage >&2
		exit 1
		;;
esac

# Magic byte sequences (hex) to identify kernel, SquashFS and bootloader
KERNEL_MAGIC="8112EDFE"
SQUASHFS_MAGIC_LE="68737173"
SQUASHFS_MAGIC_BE="73717368"
BOOTLOADER_MAGIC_LE="00908040"
BOOTLOADER_MAGIC_BE="40809000"

output_dir="$(basename "$input_file" | sed -r 's/(.*)\.[^\.]*/\1/')"
mkdir -p "$output_dir"

extract_bootloader
extract_kernel_filesystem