#!/bin/sh
#===============================================================================
#
#  recovery-initramfs-init
#
#  Copyright (C) 2016, 2017 by Digi International Inc.
#  All rights reserved.
#
#  This program is free software; you can redistribute it and/or modify it
#  under the terms of the GNU General Public License version 2 as published by
#  the Free Software Foundation.
#
#
#  !Description: Init script for recovery initramfs
#
#===============================================================================

# Variables.
#------------------------------------------------------------------------------
ENV_BOOT_RECOVERY="boot_recovery"
ENV_RECOVERY_COMMAND="recovery_command"

SW_CONFIG="/etc/swupdate.cfg"
PUBLIC_KEY="/etc/ssl/certs/key.pub"

USB_MOUNT_DIR="/run/media"
UPDATE_MOUNT_DIR="/mnt/update"

ROOTFS_IMAGE_IN_PACKAGE="no"
ENCRYPT_ROOTFS="no"
SWUPDATE_OUTPUT="swupdate_output.txt"

PART_LIST=""
ENC_PARTS=""
DEFAULT_ENC_PARTS="yes"

NAND_PARTS_BLACKLIST="bootloader environment linux recovery safe"
EMMC_PARTS_BLACKLIST="linux recovery safe"

ENC_DIFF=""
UNENC_DIFF=""

REBOOT_TIME=10

# Functions.
#------------------------------------------------------------------------------
#------------------------------------------------------------------------------
# Function - log
#
# Prints the given text in the console.
#
# @param ${1}  - Text to print.
#------------------------------------------------------------------------------
log() {
	echo "[RECOVERY] ${1}" 1>&2
}

#------------------------------------------------------------------------------
# Function - log_warning
#
# Prints the given text in the console as a warning.
#
# @param ${1}  - Warning text to print.
#------------------------------------------------------------------------------
log_warning() {
	log "[WARNING] ${1}"
}

#------------------------------------------------------------------------------
# Function - log_error
#
# Prints the given text in the console as an error.
#
# @param ${1}  - Error text to print.
#------------------------------------------------------------------------------
log_error() {
	log "[ERROR] ${1}"
	psplash_message "ERROR: ${1}"
	psplash_progress "0"
}

#------------------------------------------------------------------------------
# Function - clear_uboot_vars
#
# Clears recovery U-Boot variables.
#------------------------------------------------------------------------------
clear_uboot_vars() {
	fw_setenv "${ENV_BOOT_RECOVERY}"
	fw_setenv "${ENV_RECOVERY_COMMAND}"
}

#------------------------------------------------------------------------------
# Function - read_uboot_var
#
# Reads the given U-Boot variable.
#
# @param ${1}  - U-Boot variable to read.
# @param ${2}  - Where to store the value of the read variable.
#------------------------------------------------------------------------------
read_uboot_var() {
	eval "${2}=\"$(fw_printenv -n ${1} 2>/dev/null)\""
}

#------------------------------------------------------------------------------
# Function - set_uboot_var
#
# Sets the given U-Boot variable.
#
# @param ${1}  - U-Boot variable to set.
# @param ${2}  - Value to set.
#------------------------------------------------------------------------------
set_uboot_var() {
	fw_setenv ${1} "${2}" 2>/dev/null
}

#------------------------------------------------------------------------------
# Function - psplash_message
#
# Shows the given message in the psplash screen.
#
# @param ${1}  - Message to show.
#------------------------------------------------------------------------------
psplash_message() {
	echo "MSG ${1}" > /tmp/psplash_fifo
	sleep 0.2
}

#------------------------------------------------------------------------------
# Function - psplash_progress
#
# Sets the psplash progress bar percentage to the given one.
#
# @param ${1}  - Progress percentage.
#------------------------------------------------------------------------------
psplash_progress() {
	echo "PROGRESS ${1}" > /tmp/psplash_fifo
	sleep 0.2
}

#------------------------------------------------------------------------------
# Function - reboot_system
#
# Reboots the system.
#------------------------------------------------------------------------------
reboot_system() {
	sync && reboot -f
}

#------------------------------------------------------------------------------
# Function - quit_with_error
#
# Ends the recovery process with the given error message.
#
# @param ${1}  - Error message.
#------------------------------------------------------------------------------
quit_with_error() {
	clear_uboot_vars
	log_error "${1}"
	log "The system will now reboot in ${REBOOT_TIME} seconds"
	sleep "${REBOOT_TIME}"
	reboot_system
}

#------------------------------------------------------------------------------
# Function - is_nand
#
# Verifies if the system is running in a NAND flash.
#
# @return  - "yes" if the system is running in NAND, "no" otherwise
#------------------------------------------------------------------------------
is_nand() {
	if grep -qs "\broot=\(ubi\|mtd\)" /proc/cmdline; then
		echo "yes"
	else
		echo "no"
	fi
}

#------------------------------------------------------------------------------
# Function - contains
#
# Returns true if and only if the list $1 contains $2
#------------------------------------------------------------------------------
contains() {
	echo "${1}" | grep -qs "\b${2}\b"
	return $?
}

#------------------------------------------------------------------------------
# Function - remove_duplicates
#
# Removes duplicate entries from a list
#------------------------------------------------------------------------------
remove_duplicates () {
	echo "${1}" | tr ' ' '\n' | sort | uniq | tr '\n' ' ' | xargs
}

#------------------------------------------------------------------------------
# Function - remove_entry
#
# Removes an entry from a list
#------------------------------------------------------------------------------
remove_entry () {
	echo "${1}" | sed "s/\b${2}\b//g" | xargs
}

#------------------------------------------------------------------------------
# Function - format_partition
#
# Formats the given partition.
#
# @param ${1}  - Partition name to format.
#------------------------------------------------------------------------------
format_partition() {
	if [ "$(is_nand)" = "yes" ]; then
		format_ubi_volume "${1}"
	else
		format_emmc_block "${1}"
	fi
}

#------------------------------------------------------------------------------
# Function - format_ubi_volume
#
# Formats and re-creates the given UBI volume.
#
# @param ${1}  - UBI Volume name to format.
#------------------------------------------------------------------------------
format_ubi_volume() {
	log "Formatting '${1}' ubi volume"
	psplash_message "Formatting '${1}' partition..."
	psplash_progress "0"

	# Read the singlemtdsys variable.
	read_uboot_var singlemtdsys singlemtdsys

	if [ "${singlemtdsys}" = "yes" ]; then
		# Find the volume number associated to the volume name
		for d in /dev/ubi0_*; do
			volname="$(ubinfo ${d} | grep ^Name | awk '{print $(2)}')"
			if [ "${volname}" = "${1}" ]; then
				# Find mountpoint
				mountpoint="$(mount | grep ubi0:${1} | awk '{print $(3) }')"
				umount ${mountpoint} 2> /dev/null
				# Wipe out volume
				ubiupdatevol ${d} -t
			fi
		done
		psplash_progress "100"
	else
		# Find the MTD partition.
		local mtd_num="$(sed -ne "s/mtd\([0-9]\+\):.*\<${1}\>.*/\1/g;T;p" /proc/mtd 2>/dev/null)"
		if [ -z "${mtd_num}" ]; then
			quit_with_error "Could not find MTD partition for volume '${1}'"
		else
			# Umount in case partition is mounted, ignore errors.
			if grep -qs "${1}" /proc/mounts; then
				local path="$(sed -ne "s/.*:${1} \(.*\) ubifs.*/\1/g;T;p" /proc/mounts 2>/dev/null)"
				umount "${path}" >/dev/null 2>&1
			fi
			ubidetach -p "/dev/mtd${mtd_num}" >/dev/null 2>&1
			# Format MTD partition.
			if ! ubiformat "/dev/mtd${mtd_num}" -q -y; then
				quit_with_error "Error erasing '/dev/mtd${mtd_num}' block"
			fi
			psplash_progress "50"
			# Attach and get UBI device number
			local dev_number="$(ubiattach -p /dev/mtd${mtd_num} 2>/dev/null | sed -ne 's,.*device number \([0-9]\).*,\1,g;T;p' 2>/dev/null)"
			# Create UBI Vol.
			ubimkvol "/dev/ubi${dev_number}" -m -N "${1}" >/dev/null 2>&1
			if [ "$?" = "0" ]; then
				# Configure the empty UBIFS partition to use ZLIB
				[ "${1}" = "update" ] && UBIFS_COMPRESSION="-x zlib"

				volid="$(ubinfo "/dev/ubi${dev_number}" -N "${1}" | sed -ne 's,Volume ID:[[:blank:]]\+\([0-9]\+\)[[:blank:]]\+.*,\1,g;T;p')"
				mkfs.ubifs ${UBIFS_COMPRESSION} -F /dev/ubi${dev_number}_${volid}
				psplash_progress "100"
				log "Partition '${1}' successfully erased!"
				# Detach MTD partition.
				ubidetach -p "/dev/mtd${mtd_num}" >/dev/null 2>&1
			else
				quit_with_error "Error creating '${1}' UBI volume"
			fi
		fi
	fi
}

#------------------------------------------------------------------------------
# Function - format_emmc_block
#
# Formats the given emmc partition block name.
#
# @param ${1}  - Partition name to format.
#------------------------------------------------------------------------------
format_emmc_block() {
	log "Formatting '${1}' eMMC block"
	psplash_message "Formatting '${1}' partition..."
	psplash_progress "0"

	# If partition is encrypted, it might be open and even mounted at this point.
	local mapped_block="/dev/mapper/crypt${1}"
	if [ -e "${mapped_block}" ]; then
		# Umount in case partition is mounted, ignore errors.
		if grep -qs "${mapped_block}" /proc/mounts; then
			umount "${mapped_block}" >/dev/null 2>&1
		fi
		# Close mapped device
		cryptsetup close crypt${1}
	fi

	# Find partition block number.
	local partition_block="/dev/mmcblk0p$(fdisk -l /dev/mmcblk0 | sed -ne "s,^[^0-9]*\([0-9]\+\).*\<${1}\>.*,\1,g;T;p")"
	if [ -b "${partition_block}" ]; then
		# Umount in case partition is mounted, ignore errors.
		if grep -qs "${partition_block}" /proc/mounts; then
			umount "${partition_block}" >/dev/null 2>&1
		fi
		# If partition is encrypted, format with trustfence-tool first
		if contains "${encrypt_partitions}" "${1}"; then
			trustfence-tool --format ${partition_block} crypt${1}
			if [ ! "$?" = "0" ]; then
				quit_with_error "Error formatting '${1}' partition for encryption"
			fi
			partition_block="/dev/mapper/crypt${1}"
			psplash_progress "50"
		fi

		# In the case of the rootfs, there's no need to format the
		# partition, because we know an image will be written from
		# an update package.
		if [ "${1}" = "rootfs" ]; then
			psplash_progress "100"
			log "Partition '${1}' successfully erased!"
			return
		fi

		# Format emmc block.
		mkfs.ext4 "${partition_block}" >/dev/null 2>&1
		if [ "$?" = "0" ]; then
			psplash_progress "100"
			log "Partition '${1}' successfully erased!"
		else
			quit_with_error "Error erasing '${1}' partition"
		fi
	else
		quit_with_error "Could not find partition block for '${1}'"
	fi
}

#------------------------------------------------------------------------------
# Function - swu_package_path
#
# Get absolute path of update package searching in local media
#
# @param ${1}  - SWU package
#------------------------------------------------------------------------------
swu_package_path() {
	# Check whether the package is local and get the absolute path.
	if echo "${1}" | grep -qs '^file://'; then
		local pkg_name="$(basename ${1})"

		for i in ${UPDATE_MOUNT_DIR} $(echo ${USB_MOUNT_DIR}/*); do
			echo $i | grep -qs "${USB_MOUNT_DIR}/\*" && continue
			if [ -f "${i}/${pkg_name}" ]; then
				swu_abspath="${i}/${pkg_name}"
				break
			fi
		done

		if [ -n "${swu_abspath}" ]; then
			echo "${swu_abspath}"
		else
			quit_with_error "Unable to find update package '${pkg_name}'"
		fi
	else
		echo "${1}"
	fi
}

#------------------------------------------------------------------------------
# Function - check_swu_package
#
# Check if the update package is a valid one and get its absolute path
#
# @param ${1}  - SWU package
#------------------------------------------------------------------------------
check_swu_package() {
	log "Checking update package '$(basename ${1})'"

	if [ -z "${1}" ]; then
		quit_with_error "Firmware update package not specified"
	fi

	update_package="$(swu_package_path ${1})"

	# Check software update package.
	if [ -f "${PUBLIC_KEY}" ]; then
		swupdate -c -v -i "${update_package}" -e "${SWUPDATE_IMAGE_SET}" -k "${PUBLIC_KEY}" > "${SWUPDATE_OUTPUT}"
	else
		swupdate -c -v -e "${SWUPDATE_IMAGE_SET}" -i "${update_package}" > "${SWUPDATE_OUTPUT}"
	fi

	if [ "$?" != "0" ]; then
		quit_with_error "Invalid update package '$(basename ${1})'"
	fi

	# Check if the update package contains a rootfs image
	if [ "$(is_nand)" = "yes" ]; then
		grep "Found Image" "${SWUPDATE_OUTPUT}" | grep -qs "rootfs for handler" && ROOTFS_IMAGE_IN_PACKAGE="yes"
	else
		# For eMMC packages, check for a *.ext4 image, which is very likely intended for the rootfs partition
		grep -E "Found( compressed){0,1} Image" "${SWUPDATE_OUTPUT}" | grep -Eqs "\.ext4(\.gz){0,1} in device" && ROOTFS_IMAGE_IN_PACKAGE="yes"
	fi

	# Check if the rootfs is meant to be encrypted
	if [ "${ROOTFS_IMAGE_IN_PACKAGE}" = "yes" ]; then
		grep "Description" "${SWUPDATE_OUTPUT}" | grep -qs "Encrypted rootfs" && ENCRYPT_ROOTFS="yes"
	fi
}

#------------------------------------------------------------------------------
# Function - update_mtdparts
# Updates the encryption flags on the mtdparts environment variable
#------------------------------------------------------------------------------
update_mtdparts() {
	# Read the mtdparts variable.
	read_uboot_var mtdparts mtdparts

	# Check if there is any command.
	if [ -z "${mtdparts}" ]; then
		quit_with_error "No mtdparts found"
	fi

	# Add encryption flag to new encrypted partitions
	for p in ${ENC_DIFF}; do
		mtdparts=$(echo "${mtdparts}" | sed "s/(${p})/(${p})enc/g")
	done

	# Remove encryption flag from unencrypted partitions
	for p in ${UNENC_DIFF}; do
		mtdparts=$(echo "${mtdparts}" | sed "s/(${p})enc/(${p})/g")
	done

	set_uboot_var mtdparts "${mtdparts}"
}

#------------------------------------------------------------------------------
# Function - parse_partition_info
#
# Obtain a list of all partitions and a list of encrypted partitions
#------------------------------------------------------------------------------
parse_partition_info() {
	if [ "$(is_nand)" = "yes" ]; then
		# Read the mtdparts variable.
		read_uboot_var mtdparts mtdparts

		if [ -z "${mtdparts}" ]; then
			quit_with_error "No mtdparts found"
		fi

		TMP_LIST=$(echo ${mtdparts} | cut -d ':' -f2 | tr "," " ")
		for p in ${TMP_LIST}; do
			PART=$(echo $p | sed 's,.*(\(.*\)).*,\1,')
			PART_LIST="${PART_LIST} ${PART}"
			echo $p | grep -qs ")enc" && ENC_PARTS="${ENC_PARTS} ${PART}"
		done
	else
		PART_LIST=$(fdisk -l /dev/mmcblk0 | grep '^  *' | rev | cut -d ' ' -f1 | rev)
		# If the partition list doesn't exist, we can consider it empty
		read_uboot_var encrypted_parts_list ENC_PARTS
		ENC_PARTS=$(remove_duplicates "${ENC_PARTS}")
	fi

	# If a new partition list wasn't explicitly passed in the recovery
	# command, assume the list hasn't changed. This prevents having to
	# always pass the list even if it doesn't have any changes.
	[ "${DEFAULT_ENC_PARTS}" = "yes" ] && encrypt_partitions="${ENC_PARTS}"
}

#------------------------------------------------------------------------------
# Function - get_blacklist
#
# Return the list of partitions that can't be encrypted on the storage media
#------------------------------------------------------------------------------
get_blacklist() {
	if [ "$(is_nand)" = "yes" ]; then
		echo "${NAND_PARTS_BLACKLIST}"
	else
		echo "${EMMC_PARTS_BLACKLIST}"
	fi
}

# Main
#------------------------------------------------------------------------------
# Setup the environment.
export PATH=/bin:/sbin:/usr/bin:/usr/sbin

# Mount virtual file system.
mkdir -p /proc /sys /dev /tmp
mount -t proc proc /proc
mount -t sysfs sysfs /sys
mount -t devtmpfs devtmpfs /dev
mount -t tmpfs tmpfs /tmp

# Register mdev as device manager
echo > /dev/mdev.seq
echo > /dev/mdev.log
mdev -d
mdev -s

# Give some time for the devices to settle down so mdev can mount all of them
sleep 2

# Run all shell scripts in postinstall folder
run-parts /etc/*-postinsts

# Setup fw_printenv.
mkdir -p /var/lock

# Set kernel console loglevel.
sysctl -q -w kernel.printk=4

# Set path for psplash communication FIFO
export PSPLASH_FIFO_DIR="/tmp"

# Start psplash.
psplash &

# Parse the kernel command line.
for arg in $(cat /proc/cmdline); do
	case "${arg}" in
		rescue=1) eval "${arg}";;
	esac
done

# Jump to a rescue shell if requested.
if [ -n "${rescue}" ]; then
	# Expand console and respawn if exited
	while true; do
		setsid cttyhack sh -l
		sleep 1
	done
fi

log "Starting recovery..."
psplash_message "Starting recovery..."

# Read the recovery command.
read_uboot_var "${ENV_RECOVERY_COMMAND}" COMMAND

# Check if there is any command.
if [ -z "${COMMAND}" ]; then
	quit_with_error "No command found"
fi

# Parse the recovery command.
for arg in ${COMMAND}; do
	case "${arg}" in
		wipe_update)
			wipe_update_bool=true;;
		encryption_key=*)
			encryption_key_bool=true;
			eval "${arg}";;
		update_package=*)
			update_package_bool=true;
			eval "${arg}";;
		encrypt_partitions=*)
			eval "${arg}";
			DEFAULT_ENC_PARTS="no";
			encrypt_partitions=$(echo ${encrypt_partitions} | tr "," " ");
			encrypt_partitions=$(remove_duplicates "${encrypt_partitions}");;
		wipe_ubi_partitions=*)
			eval "${arg}";
			wipe_ubi_partitions=$(echo ${wipe_ubi_partitions} | tr "," " ");
			wipe_ubi_partitions=$(remove_duplicates "${wipe_ubi_partitions}");;
		*)
			# Not supported command
			quit_with_error "Unknown recovery command '${arg}'";;
	esac
done

# Get current partition information
parse_partition_info

# Sanitize the encrypted partition list
BLACKLISTED=""
for p in ${encrypt_partitions}; do
	contains "${PART_LIST}" "${p}" || quit_with_error "Cannot encrypt nonexistant partition '${p}'"
	# Take note of any blacklisted partitions to remove them from the list later
	contains "$(get_blacklist)" "${p}" && BLACKLISTED="${BLACKLISTED} ${p}"
done
for p in ${BLACKLISTED}; do
	log_warning "Encryption of partition '${p}' is forbidden, skipping"
	encrypt_partitions=$(remove_entry "${encrypt_partitions}" "${p}")
done

# Select update package image
if [ "$(is_nand)" = "yes" ]; then
	SWUPDATE_IMAGE_SET="mtd,platform"
else
	SWUPDATE_IMAGE_SET="mmc,platform"
fi

# On eMMC, if the 'update' partition is encrypted, we need to mount it manually
if [ "$(is_nand)" = "no" ] && contains "${ENC_PARTS}" "update"; then
	update_block="/dev/mmcblk0p$(fdisk -l /dev/mmcblk0 | sed -ne "s,^[^0-9]*\([0-9]\+\).*\<update\>.*,\1,g;T;p")"
	trustfence-tool ${update_block} cryptupdate
	if [ "$?" = "0" ]; then
		# Reset block path to the decrypted mapped device
		update_block="/dev/mapper/cryptupdate"

		mkdir -p ${UPDATE_MOUNT_DIR}
		FSTYPE="$(blkid ${update_block} | sed -e 's,.*TYPE="\([^"]\+\)".*,\1,g')"
		mount ${FSTYPE:+-t ${FSTYPE}} ${update_block} ${UPDATE_MOUNT_DIR}
	else
		quit_with_error "Error mounting encrypted update partition"
	fi
fi

# Sanity checks.
if [ -n "${update_package_bool}" ]; then
	check_swu_package "${update_package}"
fi

# Format UBI partitions
if [ "$(is_nand)" = "yes" -a -n "${wipe_ubi_partitions}" ]; then
	for p in ${wipe_ubi_partitions}; do
		# Only format 'update' partition if it doesn't have a pending
		# update package.
		[ "${p}" = "update" ] && \
		echo "${update_package}" | grep -qs "^${UPDATE_MOUNT_DIR}" && continue

		# 'rootfs' partition is formatted prior to the update, no need
		# to format it now.
		[ "${p}" = "rootfs" ] && continue

		# To protect against manually injected commands, only format
		# partitions that exist and aren't blacklisted.
		contains "${PART_LIST}" "${p}" && \
		! contains "${NAND_PARTS_BLACKLIST}" "${p}" && \
		format_ubi_volume "${p}"
	done
fi

# Compare the current and new encrypted partition lists and take note of the
# differences, if any. If we know the list hasn't changed, don't bother
# checking.
if [ "${DEFAULT_ENC_PARTS}" = "no" ]; then
	# Iterate through the old encrypted partition list
	for p in ${ENC_PARTS}; do
		if ! contains "${encrypt_partitions}" "${p}"; then
			# Partition is no longer in the list, unencrypt it
			UNENC_DIFF="${UNENC_DIFF} ${p}"
		fi
	done

	# Iterate through the new encrypted partition list
	for p in ${encrypt_partitions}; do
		if ! contains "${ENC_PARTS}" "${p}"; then
			# Partition is new in the list, encrypt it
			ENC_DIFF="${ENC_DIFF} ${p}"
		fi
	done
fi

# If we have an update package in the 'update' partition and we also need to
# format said partition, avoid doing so by aborting the operation or by
# avoiding the formatting.
if echo "${update_package}" | grep -qs "^${UPDATE_MOUNT_DIR}"; then
	# In the case of a state change (encrypted->unencrypted or
	# unencrypted->encrypted), we can simply cancel it and continue with
	# the update. Make sure to remove leading/trailing whitespaces from
	# the diffs.
	if contains "${UNENC_DIFF}" "update"; then
		log_warning "Update package in 'update' partition, skip the partition's unencryption process."

		encrypt_partitions="${encrypt_partitions} update"
		UNENC_DIFF=$(remove_entry "${UNENC_DIFF}" "update")
	elif contains "${ENC_DIFF}" "update"; then
		log_warning "Update package in 'update' partition, skip the partition's encryption process."

		encrypt_partitions=$(remove_entry "${encrypt_partitions}" "update")
		ENC_DIFF=$(remove_entry "${ENC_DIFF}" "update")
	fi

	# Worst case scenario: the partition remains encrypted, but the key is
	# going to be changed. Since this affects all encrypted partitions,
	# which might also include the rootfs, we can't cancel the key change
	# and continue with the update. To avoid unexpected behaviour, abort
	# both operations and quit.
	if [ -n "${encryption_key_bool}" ] && \
	   contains "${ENC_PARTS}" "update" && \
	   contains "${encrypt_partitions}" "update"; then
		quit_with_error "Cannot change the encryption key with an update package in an encrypted 'update' partition."
	fi
fi

# Unconditionally remove the 'rootfs' partition from the diffs, since manual
# changes aren't allowed.
if contains "${UNENC_DIFF}" "rootfs"; then
	log_warning "'rootfs' partition encryption status cannot be changed manually"

	encrypt_partitions="${encrypt_partitions} rootfs"
	UNENC_DIFF=$(remove_entry "${UNENC_DIFF}" "rootfs")
elif contains "${ENC_DIFF}" "rootfs"; then
	log_warning "'rootfs' partition encryption status cannot be changed manually"

	encrypt_partitions=$(remove_entry "${encrypt_partitions}" "rootfs")
	ENC_DIFF=$(remove_entry "${ENC_DIFF}" "rootfs")
fi

# If the rootfs is encrypted and we change the key without providing an update,
# we will wipe it out. Abort the operation.
if [ -n "${encryption_key_bool}" ] && \
   [ "${ROOTFS_IMAGE_IN_PACKAGE}" = "no" ] && \
   contains "${ENC_PARTS}" "rootfs"; then
	quit_with_error "Cannot change the encryption key with an encrypted 'rootfs' partition and no update package"
fi

# Decide if we need to (un)encrypt the rootfs based on the information provided
# by the update package.
if [ "${ROOTFS_IMAGE_IN_PACKAGE}" = "yes" ]; then
	if ! contains "${ENC_PARTS}" "rootfs" && \
	   [ "${ENCRYPT_ROOTFS}" = "yes" ]; then
		ENC_DIFF="${ENC_DIFF} rootfs"
		encrypt_partitions="${encrypt_partitions} rootfs"
	elif contains "${ENC_PARTS}" "rootfs" && \
	     [ "${ENCRYPT_ROOTFS}" = "no" ]; then
		UNENC_DIFF="${UNENC_DIFF} rootfs"
		encrypt_partitions=$(remove_entry "${encrypt_partitions}" "rootfs")
	fi
fi

# Remove leading/trailing whitespaces from the new encrypted partitions list.
encrypt_partitions=$(echo "${encrypt_partitions}" | xargs)

# Automatically program a random encryption key if there is none and we're
# about to encrypt at least one partition.
if [ -n "${ENC_DIFF}" -a -z "${encryption_key_bool}" ]; then
	trustfence-tool --key-check >/dev/null 2>&1
	RETVAL="$?"
	if [ "${RETVAL}" = "1" ]; then
		log "No encryption key detected, generating a random one"
		encryption_key_bool=true
	elif [ "${RETVAL}" != "0" ]; then
		quit_with_error "Unable to check if encryption key is set in the system"
	fi
fi

# Check if encryption key command is configured.
if [ -n "${encryption_key_bool}" ]; then
	log "Trustfence encryption key setup requested (new key: ${encryption_key:-random})"
	psplash_message "Configuring new encryption key..."
	psplash_progress "0"

	trustfence-tool "--newkey${encryption_key:+=${encryption_key}}"
	if [ "$?" = "0" ]; then
		psplash_progress "100"
		log "Trustfence encryption key setup succeed!"
	else
		quit_with_error "Error configuring trustfence encryption key"
	fi
fi

# Proceed with the formatting if any partitions require it
if [ -n "${encryption_key_bool}" -a -n "${encrypt_partitions}" ] || [ -n "${ENC_DIFF}" -o -n "${UNENC_DIFF}"  ]; then
	log "Proceeding to format partitions"

	# On NAND devices, if the list has changed, reflect the change in the
	# environment now so that the MTD driver loads the partitions with
	# their new encryption status on the second run of the recovery script.
	if [ "$(is_nand)" = "yes" ] && [ -n "${ENC_DIFF}" -o -n "${UNENC_DIFF}" ]; then
		update_mtdparts
	fi

	# Format all currently encrypted partitions with the new encryption key
	if [ -n "${encryption_key_bool}" ]; then
		ENC_DIFF="${encrypt_partitions}"
	fi

	if [ "$(is_nand)" = "yes" ]; then
		psplash_message "Preparing new recovery command..."
		psplash_progress "0"

		# Remove the encrypted parts list from the recovery command.
		new_command="${COMMAND}"
		new_command=$(echo "${new_command}" | sed "s/encrypt_partitions=[^ ]*//g")
		if [ -n "${encryption_key_bool}" ]; then
			# Modify the recovery command to not set the key again.
			new_command=$(echo "${new_command}" | sed "s/encryption_key=[^ ]*//g")
		fi

		# Append the list of partitions to format to the recovery command.
		wipe_ubi_partitions=$(echo "${ENC_DIFF} ${UNENC_DIFF}" | tr " " ",")
		new_command="${new_command} wipe_ubi_partitions=${wipe_ubi_partitions}"
		set_uboot_var "${ENV_RECOVERY_COMMAND}" "${new_command}"

		log "Rebooting to complete partition formatting..."
		psplash_progress "100"
		reboot_system
	else
		# On eMMC devices, update the new encrypted partition list
		# gradually using the old list as a base. By the end, the list
		# should be the same as the new one.
		for p in ${ENC_DIFF}; do
			format_emmc_block ${p}
			ENC_PARTS=$(remove_duplicates "${ENC_PARTS} ${p}")
			set_uboot_var encrypted_parts_list "${ENC_PARTS}"
		done

		for p in ${UNENC_DIFF}; do
			format_emmc_block ${p}
			ENC_PARTS=$(remove_entry "${ENC_PARTS}" "${p}")
			set_uboot_var encrypted_parts_list "${ENC_PARTS}"
		done
	fi
fi

# Check if update package command is configured.
if [ -n "${update_package_bool}" ]; then
	# Format the UBI volume before updating if needed.
	if [ "$(is_nand)" = "yes" ] && \
	   [ "$ROOTFS_IMAGE_IN_PACKAGE" = "yes" ]; then
		format_ubi_volume rootfs
	fi

	log "Firmware update requested"
	psplash_message "Updating firmware..."
	psplash_progress "0"

	log "Update package location: ${update_package}"
	# Execute the progress binary.
	progress -wp &
	# Execute the software update.
	if [ -f "${PUBLIC_KEY}" ]; then
		[ "$(is_nand)" = "no" ] && [ "${ENCRYPT_ROOTFS}" = "yes" ] && SWUPDATE_PREUPDATE_CMD="-P /etc/mount_cryptrootfs.sh"
		swupdate -e "${SWUPDATE_IMAGE_SET}" -f "${SW_CONFIG}" -i "${update_package}" -k "${PUBLIC_KEY}" ${SWUPDATE_PREUPDATE_CMD}
	else
		swupdate -e "${SWUPDATE_IMAGE_SET}" -f "${SW_CONFIG}" -i "${update_package}"
	fi
	if [ "$?" = "0" ]; then
		log "Firmware update process succeed!"
	else
		quit_with_error "Error executing the firmware update"
	fi
fi

# Check if wipe update patition command is configured.
if [ -n "${wipe_update_bool}" ]; then
	log "Wipe 'update' partition requested"
	format_partition update
fi

# End the recovery process.
clear_uboot_vars
reboot_system
