#!/bin/sh
#===============================================================================
#
#  update-firmware
#
#  Copyright (C) 2021-2023 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: Firmware update script
#
#===============================================================================


if [ "$(fw_printenv -n dualboot 2>/dev/null)" != "yes" ]; then
	exec update-firmware.recovery "$@"
fi

SCRIPTNAME="$(basename "$(readlink -f "${0}")")"
VERBOSE=""
PUBLIC_KEY="/etc/ssl/certs/key.pub"
ACTIVE_SYSTEM="$(fw_printenv -n active_system 2>/dev/null)"
SHOW_ACTIVE_SYSTEM=0
SCRIPT_MODE=0
REBOOT=1
UPDATE_FILE=""
ALT_BOOT=""
ALT_ROOTFS=""

# Check if the rootfs is ubifs to determine if it is a nand or emmc device
NANDROOTFS="$(grep -qs '[[:blank:]]\+/[[:blank:]]\+ubifs.*' /proc/mounts 2>/dev/null && echo 1)"

## Local functions
usage() {
	cat <<EOF

Usage: ${SCRIPTNAME} [OPTIONS] </your-path/your-filename>.swu

  -a          --active                   Show currently active system
              --no-reboot                Do not reboot after update
  -v          --verbose                  Enable verbosity
  -h          --help                     Print help and exit

EOF
}

show_active_system() {
	local act_sys="$(echo ${ACTIVE_SYSTEM} | cut -d'_' -f2)"

	if [ ${SCRIPT_MODE} -eq 0 ]; then
		act_sys="$(echo ${act_sys} | tr [:lower:] [:upper:])"
		echo "Active system is ${act_sys}"
	else
		echo "${act_sys}"
	fi
}

reboot_system() {
	if [ ${REBOOT} -eq 1 ]; then
		echo "Firmware update finished. Rebooting the system."
		reboot -f
	else
		echo "Firmware update finished. Reboot the system to use the new version."
	fi
}

swap_active_system() {
	if [ -z "${NANDROOTFS}" ]; then
		local PART_UUID=""

		# Get boot and rootfs partition index
		local MMC_PART="$(realpath /dev/disk/by-partlabel/${ALT_BOOT} | grep -o '[[:digit:]]\+$')"

		# Search rootfs UUID
		local MMCROOT_DEV="$(realpath /dev/disk/by-partlabel/${ALT_ROOTFS})"
		for uuid in /dev/disk/by-partuuid/*; do
			if [ "$(realpath "${uuid}")" = "${MMCROOT_DEV}" ]; then
				PART_UUID="$(basename "${uuid}")"
				break
			fi
		done

		if [ -z "${MMC_PART}" ] || [ -z "${PART_UUID}" ]; then
			echo "[ERROR] Unable to detect partitions."
			return 1
		fi

		fw_setenv mmcroot "PARTUUID=${PART_UUID}"
		fw_setenv mmcpart "${MMC_PART}"
	else
		fw_setenv mtdbootpart ${ALT_BOOT}
		fw_setenv mtdrootfspart ${ALT_ROOTFS}
		fw_setenv rootfsvol ${ALT_ROOTFS}
	fi

	fw_setenv active_system ${ALT_BOOT}
	fw_setenv bootcount 0

	return 0
}

update_emmc() {
	local IMAGE_SET="mmc,primary"

	# Get current partition information so we can
	# determine where to flash the images.
	[ "${ACTIVE_SYSTEM}" = "linux_a" ] && IMAGE_SET="mmc,secondary"

	echo ""
	echo "Updating '${IMAGE_SET}' image set from '${UPDATE_FILE}'..."
	echo ""

	# Execute the update.
	swupdate ${VERBOSE} -i "${UPDATE_FILE}" -e "${IMAGE_SET}"
}

update_nand() {
	local IMAGE_SET="mtd,primary"

	# Get current partition information so we can
	# determine where to flash the images.
	[ "${ACTIVE_SYSTEM}" = "linux_a" ] && IMAGE_SET="mtd,secondary"

	echo ""
	echo "Updating '${IMAGE_SET}' image set from '${UPDATE_FILE}'..."
	echo ""

	# Execute the update.
	if [ -f "${PUBLIC_KEY}" ]; then
		swupdate ${VERBOSE} -i "${UPDATE_FILE}" -e "${IMAGE_SET}" -k "${PUBLIC_KEY}"
	else
		swupdate ${VERBOSE} -i "${UPDATE_FILE}" -e "${IMAGE_SET}"
	fi
}

update_device() {
	local ret

	show_active_system
	echo "Updating system on $(echo ${ALT_BOOT} | cut -d'_' -f2 | tr [:lower:] [:upper:])"

	if [ -z "${NANDROOTFS}" ]; then
		update_emmc
	else
		update_nand
	fi

	if [ "$?" = "0" ]; then
		if ! swap_active_system; then
			exit 1
		fi
		reboot_system
	else
		echo "[ERROR] $? There was an error performing the update"
	fi
}

while :; do
	case $1 in
		-a|--active) SHOW_ACTIVE_SYSTEM=1
		;;
		--no-reboot) REBOOT=0
		;;
		-s) SCRIPT_MODE=1
		;;
		-v|--verbose) VERBOSE="-v"
		;;
		-h|--help) usage;exit
		;;
		*) UPDATE_FILE="${1}"
		    break
		;;
	esac
	shift
done

# Show active system.
if [ ${SHOW_ACTIVE_SYSTEM} -eq 1 ]; then
	show_active_system
	exit
fi

if [ "${ACTIVE_SYSTEM}" = "linux_a" ]; then
	ALT_BOOT="linux_b"
	ALT_ROOTFS="rootfs_b"
else
	ALT_BOOT="linux_a"
	ALT_ROOTFS="rootfs_a"
fi

# Check update file parameter.
if [ -z "${UPDATE_FILE}" ]; then
	echo "[ERROR] Update file not specified"
	exit
elif [ ! -f "${UPDATE_FILE}" ]; then
	echo "[ERROR] Update file '${UPDATE_FILE}' does not exist"
	exit
fi

update_device
