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

REBOOT_TIME=10

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

#------------------------------------------------------------------------------
# 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 [ -f /proc/mtd ] && grep -qs mtd /proc/mtd; then
		echo "yes"
	else
		echo "no"
	fi
}

#------------------------------------------------------------------------------
# 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() {
	# 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
			ubidetach -p "/dev/mtd${mtd_num}" >/dev/null 2>&1
		fi
		# 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
}

#------------------------------------------------------------------------------
# Function - format_emmc_block
#
# Formats the given emmc partition block name.
#
# @param ${1}  - Partition name to format.
#------------------------------------------------------------------------------
format_emmc_block() {
	# Find partition block number.
	local partition_block="/dev/mmcblk0p$(parted -s /dev/mmcblk0 print | 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
		# 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 name
#------------------------------------------------------------------------------
swu_package_path() {
	for i in ${UPDATE_MOUNT_DIR} $(echo ${USB_MOUNT_DIR}/*); do
		echo $i | grep -qs "${USB_MOUNT_DIR}/\*" && continue
		if [ -f "${i}/${1}" ]; then
			swu_abspath="${i}/${1}"
			break
		fi
	done

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

#------------------------------------------------------------------------------
# Function - set_encryption_flag
#
# Set the rootfs encryption flag to the mtdparts variable.
#------------------------------------------------------------------------------
set_encryption_flag() {
	if [ "$(is_nand)" = "no" ]; then
		return
	fi

	# 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

	# Parse the mtdparts value.
	case "${mtdparts}" in
		*\(rootfs\)enc*)
			# Partition already flagged.
			;;
		*\(rootfs\)*)
			# Add the flag to the rootfs.
			local new_mtdparts=$(echo "${mtdparts}" | sed "s/(rootfs)/(rootfs)enc/g")
			set_uboot_var mtdparts "${new_mtdparts}"
			;;
		*)
			quit_with_error "Error flagging rootfs as encrypted"
			;;
	esac

	if [ -n "${update_package_bool}" ]; then
		# Modify the recovery command not to set again the key.
		new_command=$(echo "${COMMAND}" | sed "s/encryption_key=[^ ]*//g")
		set_uboot_var "${ENV_RECOVERY_COMMAND}" "${new_command}"
		psplash_progress "100"
		reboot_system
	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
if [ -f /proc/sys/kernel/hotplug ]; then
	echo > /dev/mdev.seq
	echo > /dev/mdev.log
	echo "/sbin/mdev" > /proc/sys/kernel/hotplug
	mdev -s
fi

# Give some time for the devices to settle down
sleep 2

# Setup fw_printenv.
mkdir -p /var/lock

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

if [ "$(is_nand)" = "no" ]; then
	# Launch 'rngd' to feed random data to kernel entropy pool
	mkdir -p /var/run && rngd
fi

# 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}";;
		*)
			# Not supported command
			quit_with_error "Unknown recovery command '${arg}'";;
	esac
done

# 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 "10"
		log "Trustfence encryption key setup succeed!"
	else
		quit_with_error "Error configuring trustfence encryption key"
	fi

	# Set the encryption flag to the rootfs.
	set_encryption_flag

	# Format partition.
	if [ "$(is_nand)" = "no" ]; then
		psplash_message "Formatting rootfs partition..."
		rootfs_block="/dev/mmcblk0p$(parted -s /dev/mmcblk0 print | sed -ne "s,^[^0-9]*\([0-9]\+\).*\<rootfs\>.*,\1,g;T;p")"
		trustfence-tool --format ${rootfs_block} cryptroot
	fi
	psplash_progress "100"
fi

# Check if update package command is configured.
if [ -n "${update_package_bool}" ]; then
	log "Firmware update requested"
	psplash_message "Updating firmware..."
	psplash_progress "0"
	if [ -z "${update_package}" ]; then
		quit_with_error "Firmware update package not specified"
	else
		# Format the UBI volume before updating.
		if [ "$(is_nand)" = "yes" ]; then
			psplash_message "Formatting rootfs partition..."
			format_ubi_volume rootfs
			psplash_progress "0"
		fi

		# Check whether the package is local and get the absolute path
		if echo "${update_package}" | grep -qs '^file://'; then
			update_package="$(swu_package_path $(basename ${update_package}))"
		fi

		log "Update package location: ${update_package}"
		# Execute the progress binary.
		progress -wp &
		# Execute the software update.
		if [ -f "${PUBLIC_KEY}" ]; then
			swupdate -f "${SW_CONFIG}" -i "${update_package}" -k "${PUBLIC_KEY}"
		else
			swupdate -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
fi

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

# End the recovery process.
clear_uboot_vars
reboot_system
