#! /bin/bash
# -----------------------------------------------------------------------------
# Script Name: spice3_postinst
# Description: This script handles the post-installation process for updating
#              the preloader and bootloader on the system. It ensures that the
#              updates are applied correctly, logs all actions with timestamps,
#              and reboots the system if necessary.
#              The calling init-script will be removed from the runlevel after
#              execution.
# Called By:   This script is executed by the init-script /etc/init.d/spice3_postinst
# Version:     1.0
# -----------------------------------------------------------------------------
# Usage:
#   Run this script as part of the post-installation process. It will:
#     - Update the preloader if necessary.
#     - Update the bootloader if necessary.
#     - Log all actions to a logfile.
#     - Reboot the system if updates were applied.
# -----------------------------------------------------------------------------
# Notes:
#   - Ensure that the required files (preloader and bootloader) are present
#     in the specified directories before running this script.
#   - This script must be executed with root privileges.
# -----------------------------------------------------------------------------
# Exit Codes:
#   0 - Success, no reboot required.
#   2 - Error occurred during the update process.
# -----------------------------------------------------------------------------

PRELOADER_EXPECTEDSIZE="262144"

UPDATEFILEDIR="/boot/"
PRELOADER="preloader-mkpimage.bin"
BOOTLOADER="u-boot.img"
PRELOADER_SOURCE="$UPDATEFILEDIR/$PRELOADER"
PRELOADER_DEST="/dev/mmcblk0p2"
BOOTLOADER_SOURCE="$UPDATEFILEDIR/$BOOTLOADER"
BOOTLOADER_DEST="/media/fat/$BOOTLOADER"

LOGPATH="/home/spice3/log"
if [ ! -d "$LOGPATH" ]; then
	LOGPATH="/var/log"
fi
LOGFILE="$LOGPATH/spice3_postinst-$(date +%Y-%m-%d.%H%M%S).log"

RC_SUCCESS=0
RC_NEEDS_REBOOT=1
RC_ERROR=2

# -----------------------------------------------------------------------------
# Logs a message with a timestamp to the logfile.
# Arguments:
#   $1 - The message to log.
# -----------------------------------------------------------------------------
log_with_timestamp() {
	echo "$(date '+%Y-%m-%dT%H:%M:%S+%3N') $1" >>"$LOGFILE"
}

# -----------------------------------------------------------------------------
# Logs a message with a timestamp to the logfile and prints it to the console.
# Arguments:
#   $1 - The message to log and print.
# -----------------------------------------------------------------------------
log_and_print_with_timestamp() {
	echo "$(date '+%Y-%m-%dT%H:%M:%S+%3N') $1" >>"$LOGFILE"
	echo "$1"
}

# -----------------------------------------------------------------------------
# Executes a command, logs its output with timestamps, and returns the
# command's exit code.
# Arguments:
#   $@ - The command to execute.
# -----------------------------------------------------------------------------
log_command_with_timestamp() {
	echo -n "$(date '+%Y-%m-%dT%H:%M:%S+%3N') " >>"$LOGFILE"
	echo "$@" >>"$LOGFILE"
	"$@" 2>&1 | while IFS= read -r line; do
		echo "$(date '+%Y-%m-%dT%H:%M:%S+%3N') --> $line" >>"$LOGFILE"
	done
	return "${PIPESTATUS[0]}" # Return the exit code of the command
}

# -----------------------------------------------------------------------------
# Updates the preloader if necessary by comparing the source and destination.
# Logs the process and handles errors or mismatches.
# Returns:
#   RC_SUCCESS if no update is needed or successful.
#   RC_NEEDS_REBOOT if an update was performed.
#   RC_ERROR if an error occurred.
# -----------------------------------------------------------------------------
update_preloader() {
	if [ ! -f "$PRELOADER_SOURCE" ]; then
		log_with_timestamp "New preloader does not exist! Skipping!"
	elif [ ! -b "$PRELOADER_DEST" ]; then
		log_with_timestamp "Preloader destination does not exist! Skipping!"
	else
		PRELOADER_SIZE=$(stat -c%s "$PRELOADER_SOURCE")
		if [ "$PRELOADER_SIZE" -ne "$PRELOADER_EXPECTEDSIZE" ]; then
			log_with_timestamp "Unexpected preloader size!"
		else
			log_command_with_timestamp cmp --bytes="$PRELOADER_SIZE" "$PRELOADER_SOURCE" "$PRELOADER_DEST"
			CMP_RESULT=$?
			if [[ $CMP_RESULT == 0 ]]; then
				log_and_print_with_timestamp "Preloader is up-to-date!"
			elif [[ $CMP_RESULT == 1 ]]; then
				log_and_print_with_timestamp "Files are different. Updating preloader!"
				log_command_with_timestamp dd if="$PRELOADER_DEST" of="$PRELOADER_SOURCE.bkup" bs=1 count="$PRELOADER_SIZE"
				log_command_with_timestamp dd if="$PRELOADER_SOURCE" of="$PRELOADER_DEST"
				CMP_RESULT=$?
				if [[ $CMP_RESULT != 0 ]]; then
					log_with_timestamp "Error during writing preloader! Trying to revert."
					log_command_with_timestamp dd if="$PRELOADER_SOURCE.bkup" of="$PRELOADER_DEST"
					CMP_RESULT=$?
					if [[ $CMP_RESULT != 0 ]]; then
						log_with_timestamp "Error reverting preloader!"
					else
						log_with_timestamp "Reverting preloader successful!"
					fi
					return $RC_ERROR
				else
					log_with_timestamp "Updating preloader successful!"
					return $RC_NEEDS_REBOOT
				fi
			else
				log_with_timestamp "Error comparing preloaders!"
			fi
		fi
	fi
	return $RC_SUCCESS
}

# -----------------------------------------------------------------------------
# Updates the bootloader if necessary by comparing the source and destination.
# Logs the process and handles errors or mismatches.
# Returns:
#   RC_SUCCESS if no update is needed or successful.
#   RC_NEEDS_REBOOT if an update was performed.
#   RC_ERROR if an error occurred.
# -----------------------------------------------------------------------------
update_bootloader() {
	if [ ! -f "$BOOTLOADER_SOURCE" ]; then
		log_with_timestamp "New bootloader does not exist! Skipping!"
	elif [ ! -f "$BOOTLOADER_DEST" ]; then
		log_with_timestamp "Bootloader destination does not exist! Skipping!"
	else
		BOOTLOADER_SIZE=$(stat -c%s "$BOOTLOADER_SOURCE")
		if [ "$BOOTLOADER_SIZE" -le "0" ]; then
			log_with_timestamp "Unexpected bootloader size!"
		else
			log_command_with_timestamp cmp "$BOOTLOADER_SOURCE" "$BOOTLOADER_DEST"
			CMP_RESULT=$?
			if [[ $CMP_RESULT == 0 ]]; then
				log_and_print_with_timestamp "Bootloader is up-to-date!"
			elif [[ $CMP_RESULT != 0 ]]; then
				log_and_print_with_timestamp "Files are different. Updating bootloader!"
				log_command_with_timestamp cp -f -v "$BOOTLOADER_DEST" "$BOOTLOADER_SOURCE.bkup"
				log_command_with_timestamp cp -f -v "$BOOTLOADER_SOURCE" "$BOOTLOADER_DEST"
				CMP_RESULT=$?
				if [[ $CMP_RESULT != 0 ]]; then
					log_with_timestamp "Error during writing bootloader! Trying to revert."
					log_command_with_timestamp cp -f -v "$BOOTLOADER_SOURCE.bkup" "$BOOTLOADER_DEST"
					CMP_RESULT=$?
					if [[ $CMP_RESULT != 0 ]]; then
						log_with_timestamp "Error reverting bootloader!"
					else
						log_with_timestamp "Reverting bootloader successful!"
					fi
					return $RC_ERROR
				else
					log_with_timestamp "Updating bootloader successful!"
					return $RC_NEEDS_REBOOT
				fi
			else
				log_with_timestamp "Error comparing bootloaders!"
			fi
		fi
	fi
	return $RC_SUCCESS
}

# -----------------------------------------------------------------------------
# Removes the `spice3_postinst` script from the system's runlevel configuration.
# Logs the process.
# -----------------------------------------------------------------------------
remove_rcsd_link() {
	log_with_timestamp "Removing spice3_postinst from runlevel."
	if [ -n "$(which update-rc.d)" ]; then
		log_command_with_timestamp update-rc.d -f spice3_postinst remove
	fi
}

# -----------------------------------------------------------------------------
# Main script execution starts here.
# -----------------------------------------------------------------------------
log_with_timestamp "Starting spice3_postinst after firmware update..."

update_preloader
RETURN_PRE=$?
if [[ $RETURN_PRE == "$RC_ERROR" ]]; then
	log_and_print_with_timestamp "Update failed! Unable to update preloader!"
	exit $RC_ERROR
fi
update_bootloader
RETURN_BOOT=$?
if [[ $RETURN_BOOT == "$RC_ERROR" ]]; then
	log_and_print_with_timestamp "Update failed! Unable to update bootloader!"
	exit $RC_ERROR
fi

remove_rcsd_link

log_with_timestamp "Finished spice3_postinst after firmware update..."

if [[ $RETURN_PRE == "$RC_NEEDS_REBOOT" || $RETURN_BOOT == "$RC_NEEDS_REBOOT" ]]; then
	log_and_print_with_timestamp "Rebooting one time after firmware update."
	reboot
fi

exit $RC_SUCCESS
