#!/bin/bash

#
# Uncompress LZMA-packed firmware kernel.
#
# Based on the format description available at http://web.archive.org/20200701000000/www.wehavemorefun.de/fritzbox/LZMA-Kernel
#

SELF=$(readlink -f ${0})

usage() {
cat << EOF
Usage: $(basename $SELF) [-h|--help] [-u|--unpack] [-d|--dump] [-i|--info] <input-file> [output-file]

    -h|--help        print this help and exit

    -u|--unpack      unpack contained LZMA-record(s), default
    -d|--dump        dump contained LZMA-record(s) without unpacking
    -i|--info        print meta information without writing any file

    input-file       file in EVA-kernel-format, i.e. kernel.image or kernel.raw
    output-file      name of the file output to be written to, optional
                     if omitted automatically derived from the name of the input-file
EOF
}

# process command line parameters
ARGS=$(getopt -o hudi --long help,unpack,dump,info -n "${SELF}" -- "$@")
[ $? -eq 0 ] || { usage >&2; exit 1; }
eval set -- "${ARGS}"

while true; do
	case "$1" in
		-h|--help)
			usage
			exit 0
			;;
		-u|--unpack)
			METHOD=unpack
			shift
			;;
		-d|--dump)
			METHOD=dump
			shift
			;;
		-i|--info)
			METHOD=info
			shift
			;;
		--)
			shift
			break
			;;
		*)
			echo >&2 "ERROR: internal error!"
			exit 1
			;;
	esac
done
METHOD=${METHOD:-unpack} # default to unpack if unset

if [ $# -eq 1 ] || [ "$METHOD" != "info" -a $# -eq 2 ]; then
	INPUT_FILE="$1"
	[ -e "$INPUT_FILE" ] || { echo >&2 "ERROR: \"$INPUT_FILE\" could not be read"; exit 1; }

	OUTPUT_FILE="$2"
	if [ -z "$OUTPUT_FILE" ]; then
		# no OUTPUT_FILE provided => derive OUTPUT_FILE from INPUT_FILE
		if [ "$METHOD" == "unpack" ]; then
			OUTPUT_FILE="${INPUT_FILE}.unpacked"
		elif [ "$METHOD" == "dump" ]; then
			OUTPUT_FILE="${INPUT_FILE}.lzma-compressed-dump"
		fi
	fi
else
	usage >&2; exit 1;
fi

if [ "$METHOD" == "unpack" ]; then
	declare -x UNLZMA
	UNLZMA="${UNLZMA:-$(dirname $SELF)/unlzma}"
	[ -x "$UNLZMA" ] || { echo >&2 "ERROR: this script requires unlzma tool which is not found on your system, expected location of the tool is \"$UNLZMA\""; exit 1; }
fi

declare -x CKSUM
CKSUM="${CKSUM:-$(dirname $SELF)/cksum}"

declare -x SUM
SUM="${SUM:-$(dirname $SELF)/sum}"

source "$(dirname $SELF)/freetz_bin_functions"

###

print08X() {
	printf '0x%08X\n' "$@"
}


#
# $1 - magic sequence (in little-endian notation)
# $2 - magic sequence name (used in error messages)
# $3 - max expected number of matches (unlimited if omitted)
# $4 - tight lower bound of the match, i.e. the returned offset
#      must fulfill the condition "tight-lower-bound < offset"
#      (no lower bound restriction if omitted)
# $5 - optional usually heuristic based function intended to reduce
#      the number of potential matches by doing some additional checks
#
getDecOffsetOf1stMagicSequenceMatch() {
	local sequenceNamePrefix=${2}${2:+ } # sequence name with space added
	local maxNumberOfMatches=$3          # unlimited if omitted
	local lowerBound=${4:--1}            # -1        if omitted
	local deeperCheckFct=$5              # deeper check function

	declare -a offsetCandidates=($(getDecOffsetsOfAllMatches "$INPUT_FILE" bin $(echo -n $1 | invertEndianness)))

	if [ -n "${deeperCheckFct}" ]; then
		declare -a filteredOffsetCandidates
		for ((i=0; i<${#offsetCandidates[@]}; i++)); do
			if eval "${deeperCheckFct} ${offsetCandidates[i]}"; then
				filteredOffsetCandidates+=(${offsetCandidates[i]})
			fi
		done
		offsetCandidates=("${filteredOffsetCandidates[@]}")
	fi

	for ((i=0; i<${#offsetCandidates[@]}; i++)); do
		if [ $((lowerBound)) -lt ${offsetCandidates[i]} ]; then
			local numberOfMatchesFound=$((${#offsetCandidates[@]} - i))
			if [ -n "$maxNumberOfMatches" ] && [ ! $numberOfMatchesFound -le $maxNumberOfMatches ]; then
				echo >&2 "ERROR: number of ${sequenceNamePrefix}magic sequence (0x$1) matches ($numberOfMatchesFound) exceeds the expected limit ($maxNumberOfMatches)"
				return 1
			fi

			# match fulfilling all restrictions found
			echo -n ${offsetCandidates[i]}
			return 0
		fi
	done

	local fullfiling=${4:+(fulfilling the restriction $(print08X $lowerBound) < offset) }
	echo >&2 "ERROR: no ${sequenceNamePrefix}magic sequence (0x$1) ${fullfiling}found"
	return 1
}


#
# $1 - LZMA-record offset
# $2 - TI-record name
# $3 - name of the file output to be written to
#
# returns the error code of unlzma or output redirection (depending on the $METHOD)
#
processLzmaRecord() {
	local lzmaRecordOffset=$1
	local tiRecordPrefix=$(echo -n $2 | tr -c '_a-zA-Z0-9' '_')
	tiRecordPrefix=${tiRecordPrefix}${tiRecordPrefix:+.}
	local outputFile="$3"

	local lzmaCompressedLen=$(getLEu32AtOffset      "$INPUT_FILE" $((lzmaRecordOffset +  4)))
	local lzmaUncompressedLen=$(getLEu32AtOffset    "$INPUT_FILE" $((lzmaRecordOffset +  8)))
	local lzmaCompressedChecksum=$(getLEu32AtOffset "$INPUT_FILE" $((lzmaRecordOffset + 12)))

	echo "${tiRecordPrefix}LzmaCompressedStreamLength=$(print08X $lzmaCompressedLen)"
	echo "${tiRecordPrefix}LzmaUncompressedStreamLength=$(print08X $lzmaUncompressedLen)"
	echo "${tiRecordPrefix}LzmaCompressedStreamChecksum=$(print08X $lzmaCompressedChecksum)"

	local lzmaStreamOffset=$((lzmaRecordOffset + 16))

	# verify checksum
	local lzmaCompressedChecksumComputed=$(tail -c "+$(((lzmaStreamOffset + 1) + 8))" "$INPUT_FILE" 2>/dev/null | head -c $lzmaCompressedLen 2>/dev/null | $CKSUM -lp 2>/dev/null | cut -d ' ' -f 1)
	if [ $((lzmaCompressedChecksumComputed)) -ne $((lzmaCompressedChecksum)) ]; then
		echo >&2 "WARNING: calculated LZMA-stream checksum ($(print08X $lzmaCompressedChecksumComputed)) does not match the saved one ($(print08X $lzmaCompressedChecksum))"
	fi

	if [ "$METHOD" == "info" ]; then
		return 0
	fi

	local lzmaUncompressedLenOffset=$((lzmaRecordOffset + 8))
	{
		# LZMA properties (1 byte) + LZMA dictionary size (4 bytes)
		dd if="$INPUT_FILE" bs=1 skip=$lzmaStreamOffset          count=5 2>/dev/null
		# uncompressed length
		dd if="$INPUT_FILE" bs=1 skip=$lzmaUncompressedLenOffset count=4 2>/dev/null
		# padding
		dd if=/dev/zero    bs=1                                  count=4 2>/dev/null
		# compressed data
		#   +1 = tail expects number of bytes, whereas lzmaStreamOffset is zero-based
		#   +8 = skip LZMA properties (1 byte) + LZMA dictionary size (4 bytes) + padding (3 bytes)
		tail -c "+$(((lzmaStreamOffset + 1) + 8))" "$INPUT_FILE" 2>/dev/null | head -c $lzmaCompressedLen 2>/dev/null
	} | {
		if [ "$METHOD" == "unpack" ]; then
			"$UNLZMA" > "$outputFile" 2>/dev/null
		elif [ "$METHOD" == "dump" ]; then
			cat -     > "$outputFile" 2>/dev/null
		fi
	}
}

#
# $1 - bzImage offset
# $2 - bzImage length
# $3 - name of the file output to be written to
#
processX86bzImage() {
	local bzImageOffset=$1
	local bzImageLen=$2
	local outputFile="$3"

	if [ "$METHOD" == "info" ]; then
		return 0
	fi

	# in case of bzImage we do not distinguish between "unpack" and "dump" methods
	tail -c "+$((bzImageOffset + 1))" "$INPUT_FILE" 2>/dev/null | head -c $bzImageLen >"$outputFile" 2>/dev/null
}

getTiRecordLen() {
	local tiRecordOffset=$1
	getLEu32AtOffset "$INPUT_FILE" $((tiRecordOffset + 4))
}

getLoadAddr() {
	local tiRecordOffset=$1
	getLEu32AtOffset "$INPUT_FILE" $((tiRecordOffset + 8))
}

getEntryAddr() {
	local tiRecordOffset=$1
	local tiRecordLen=$(getTiRecordLen $tiRecordOffset)

	local tiRecordEntryZero=$(getLEu32AtOffset "$INPUT_FILE" $((tiRecordOffset + tiRecordLen  + 16)))
	if [ "$tiRecordEntryZero" -ne 0 ]; then
		echo >&2 "WARNING: entry address is expected to be prefixed with one 32-bit zero word"
	fi

	getLEu32AtOffset "$INPUT_FILE" $((tiRecordOffset + tiRecordLen  + 20))
}

getTiRecordChecksum() {
	local tiRecordOffset=$1
	local tiRecordLen=$(getTiRecordLen $tiRecordOffset)
	getLEu32AtOffset "$INPUT_FILE" $((tiRecordOffset + tiRecordLen  + 12))
}

maybeTiRecord() {
	local tiRecordOffset=$1      # a _potential_ tiRecordOffset
	local tiRecordPayloadOffset
	local payloadFirst16Bytes
	local TI_AR7_MAGIC_BE=$(echo -n "$TI_AR7_MAGIC" | invertEndianness)

	tiRecordPayloadOffset=$((tiRecordOffset+12))
	payloadFirst16Bytes=$(getHexContentAtOffset "$INPUT_FILE" ${tiRecordPayloadOffset} 16 2>/dev/null) || return 1

	# heuristic, check if payload contains one of the known magic sequences
	[ "${payloadFirst16Bytes:0:${#EVA_LZMA_RECORD_MAGIC_BE}}" == "$EVA_LZMA_RECORD_MAGIC_BE" ] \
	|| \
	[ "${payloadFirst16Bytes:0:${#X86_BOOT_SECTOR_MAGIC_BE}}" == "$X86_BOOT_SECTOR_MAGIC_BE" ] \
	|| \
	[ "${payloadFirst16Bytes:0:${#TI_AR7_MAGIC_BE}}"          == "$TI_AR7_MAGIC_BE" ]
}


#
# $1 - magic sequence (in little-endian notation)
# $2 - TI-record name
# $3 - name of the file output to be written to
# $4 - tight lower bound of the magic sequence match, optional
#
# returns offset of the magic sequence by setting the global variable TI_RECORD_OFFSET
#
findAndProcessTiRecord() {
	local magic="$1"
	local tiRecordName="$2"
	local tiRecordPrefix=$(echo -n "$tiRecordName" | tr -c '_a-zA-Z0-9' '_')
	tiRecordPrefix=${tiRecordPrefix}${tiRecordPrefix:+.}
	local outputFile="$3"
	local lowerBound="$4"

	# unset return value
	TI_RECORD_OFFSET=""

	local tiRecordOffset=""
	{ tiRecordOffset=$(getDecOffsetOf1stMagicSequenceMatch $magic "$tiRecordName" 1 "$lowerBound" "maybeTiRecord"); } || return 1

	local tiRecordLoadAddr=$(getLoadAddr $tiRecordOffset)
	echo "${tiRecordPrefix}LoadAddress=$(print08X $tiRecordLoadAddr)"
	local tiRecordEntryAddr=$(getEntryAddr $tiRecordOffset)
	echo "${tiRecordPrefix}EntryAddress=$(print08X $tiRecordEntryAddr)"
	local tiRecordLen=$(getTiRecordLen $tiRecordOffset)
	echo "${tiRecordPrefix}RecordLength=$(print08X $tiRecordLen)"
	local tiRecordChecksum=$(getTiRecordChecksum $tiRecordOffset)
	echo "${tiRecordPrefix}RecordChecksum=$(print08X $tiRecordChecksum)"

	local tiRecordPayloadOffset=$((tiRecordOffset+12))

	# verify TI-record checksum
	local tiRecordPayloadSum=$(tail -c "+$((tiRecordPayloadOffset + 1))" "$INPUT_FILE" 2>/dev/null | head -c $tiRecordLen 2>/dev/null | $SUM -p 2>/dev/null | cut -d ' ' -f 1)
	local tiRecordChecksumComputed=$(( (((tiRecordLen + tiRecordLoadAddr + tiRecordPayloadSum) ^ 0xFFFFFFFF) & 0xFFFFFFFF) + 1 ))
	if [ $((tiRecordChecksumComputed)) -ne $((tiRecordChecksum)) ]; then
		echo >&2 "WARNING: calculated TI-record checksum ($(print08X $tiRecordChecksumComputed)) does not match the saved one ($(print08X $tiRecordChecksum))"
	fi

	local payloadFirst16Bytes=$(getHexContentAtOffset "$INPUT_FILE" ${tiRecordPayloadOffset} 16)
	if [ "${payloadFirst16Bytes:0:${#EVA_LZMA_RECORD_MAGIC_BE}}" == "$EVA_LZMA_RECORD_MAGIC_BE" ]; then
		processLzmaRecord $tiRecordPayloadOffset "$tiRecordName" "$outputFile"
	elif [ "${payloadFirst16Bytes:0:${#X86_BOOT_SECTOR_MAGIC_BE}}" == "$X86_BOOT_SECTOR_MAGIC_BE" ]; then
		processX86bzImage $tiRecordPayloadOffset $tiRecordLen "$outputFile"
	else
		echo >&2 "ERROR: unsupported payload within ${tiRecordName}-record (${payloadFirst16Bytes})"; return 1;
	fi
	[ $? -ne 0 ] && { echo >&2 "ERROR: failed to ${METHOD} kernel"; return 1; }

	# set return value
	TI_RECORD_OFFSET=$tiRecordOffset
}

##

TI_AR7_MAGIC=FEED1281
TI_PUMA6_MAGIC=FEED8112
DUAL_KERNEL_MAGIC=FEED9112
TI_AR7_2ND_MAGIC=FEEDB007

EVA_LZMA_RECORD_MAGIC_BE=$(echo -n 075A0201 | invertEndianness)
X86_BOOT_SECTOR_MAGIC_BE=EA0500C0078CC88ED88EC08ED031E4FB

# 1st kernel (PUMA6 boxes, x86 kernel)
if getDecOffsetOf1stMagicSequenceMatch $TI_PUMA6_MAGIC "TI-PUMA6" 1 "" "maybeTiRecord" >/dev/null 2>&1; then
	findAndProcessTiRecord $TI_PUMA6_MAGIC "TI-PUMA6" "$OUTPUT_FILE" || exit 1
	exit 0
fi

# 1st kernel (all boxes)
findAndProcessTiRecord $TI_AR7_MAGIC "TI-AR7" "$OUTPUT_FILE" || exit 1
feed1281_Offset=$TI_RECORD_OFFSET

# 2nd kernel (GRX5 boxes only)
feed9112_Offset=$(getDecOffsetOf1stMagicSequenceMatch $DUAL_KERNEL_MAGIC "DUAL-kernel" 1 "" "maybeTiRecord" 2>/dev/null)
if [ -n "$feed9112_Offset" ]; then
	if [ $((feed1281_Offset - feed9112_Offset)) -ne 12 ]; then
		echo >&2 "ERROR: DUAL-kernel magic sequence (0x$DUAL_KERNEL_MAGIC) is expected to be found exactly 12 bytes before TI-AR7 record"; exit 1
	fi

	feed1281_LoadAddr=$(getLoadAddr $feed1281_Offset)
	feed9112_LoadAddr=$(getLoadAddr $feed9112_Offset)
	if [ "$feed1281_LoadAddr" != "$feed9112_LoadAddr" ]; then
		echo >&2 "WARNING: load addresses stored in TI-AR7 record ($(print08X $feed1281_LoadAddr)) and DUAL-kernel record ($(print08X $feed9112_LoadAddr)) do not match"
	fi

	feed1281_EntryAddr=$(getEntryAddr $feed1281_Offset)
	feed9112_EntryAddr=$(getEntryAddr $feed9112_Offset)
	if [ "$feed1281_EntryAddr" != "$feed9112_EntryAddr" ]; then
		echo >&2 "WARNING: entry addresses stored in TI-AR7 record ($(print08X $feed1281_EntryAddr)) and DUAL-kernel record ($(print08X $feed9112_EntryAddr)) do not match"
	fi

	feedb007_LowerBound=$((feed1281_Offset + $(getTiRecordLen $feed1281_Offset)  + 24))
	feedb007_LowerBound=$((feedb007_LowerBound + (4 - feedb007_LowerBound%4)%4)) # pad to a 4-byte boundary

	findAndProcessTiRecord $TI_AR7_2ND_MAGIC "TI-AR7-2ND" "${OUTPUT_FILE}.2ND" $((feedb007_LowerBound-1)) || exit 1
	feedb007_Offset=$TI_RECORD_OFFSET
fi