#!/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=""
SHOW_ACTIVE_SYSTEM=0
SCRIPT_MODE=0
SWAP_ACTIVE_SYSTEM=0
REBOOT=1
UPDATE_FILE=""
ALT_BOOT=""
ALT_ROOTFS=""

# Check in the command line if root=PARTUUID to determine if the system is nand or emmc
# eMMC ==> root=PARTUUID
# NAND ==> root=ubiX:rootfs_X
# NAND (squashsf) ==> root=/dev/ubiblock0_X
EMMCROOTFS="$(grep -qs 'root=PARTUUID.*' /proc/cmdline 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
              --swap-active-system       Swap active system block.
                                         This option reboots the system, unless '--no-reboot' is specified.
  -v          --verbose                  Enable verbosity
  -h          --help                     Print help and exit

EOF
}

get_active_system() {
	if [ -z "${EMMCROOTFS}" ]; then
		# For a read-only filesystem this will be /dev/ubiblock0_X
		# For an ubifs filesystem this will be ubiX:rootfs_X
		ACTIVE_SYSTEM="$(sed -e 's/^.*root=\([^ ]*\) .*$/\1/' /proc/cmdline 2>/dev/null)"
		if ! echo "${ACTIVE_SYSTEM}" | grep -qs rootfs; then
			# From /dev/ubiblock0_X to /dev/ubi0_X
			ACTIVE_SYSTEM="/dev/ubi${ACTIVE_SYSTEM#/dev/ubiblock}"
			#Volume ID:   5 (on ubi0)
			#Type:        dynamic
			#Alignment:   1
			#Size:        1817 LEBs (230715392 bytes, 220.0 MiB)
			#State:       OK
			#Name:        rootfs_b
			#Character device major/minor: 242:6
			ACTIVE_SYSTEM="$(ubinfo "${ACTIVE_SYSTEM}" | sed -ne '/^Name/s,.* \([^[:blank:]]\+\)$,\1,g;T;p')"
		fi
	else
		local MMCROOT_DEV

		MMCROOT_DEV="$(stat -c%D /)"

		for label in /dev/disk/by-partlabel/*; do
			if [ "$(stat -c"%02t%02T" "$(realpath "${label}")")" = "${MMCROOT_DEV}" ]; then
				ACTIVE_SYSTEM="$(basename "${label}")"
				break
			fi
		done
	fi

	if [ -z "${ACTIVE_SYSTEM}" ]; then
		echo "[ERROR] Unable to get active system."
		return 1
	fi

	ACTIVE_SYSTEM="${ACTIVE_SYSTEM##*_}"

	return 0
}

show_active_system() {
	if [ ${SCRIPT_MODE} -eq 0 ]; then
		local act_sys

		act_sys="$(echo ${ACTIVE_SYSTEM} | tr [:lower:] [:upper:])"
		echo "Active system is ${act_sys}"
	else
		echo "${ACTIVE_SYSTEM}"
	fi
}

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

swap_active_system() {
	# Sanity check: Some firmware updates might request not to swap
	#               active bank after the update.
	local SWAP_BANK=$(fw_printenv -n swap_bank)
	if [ "${SWAP_BANK}" = "false" ]; then
		echo "[WARNING] Active system swap cancelled by update process."
		# Reset variable and return without swapping.
		fw_setenv swap_bank
		return 0
	fi

	if [ -z "${EMMCROOTFS}" ]; then
		fw_setenv mtdbootpart ${ALT_BOOT}
		fw_setenv mtdrootfspart ${ALT_ROOTFS}
		fw_setenv rootfsvol ${ALT_ROOTFS}
	else
		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}"

	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}" = "a" ] && IMAGE_SET="mmc,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_nand() {
	local IMAGE_SET="mtd,primary"

	# Get current partition information so we can
	# determine where to flash the images.
	[ "${ACTIVE_SYSTEM}" = "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 [ -n "${EMMCROOTFS}" ]; then
		update_emmc
	else
		update_nand
	fi

	if [ "$?" = "0" ]; then
		if ! swap_active_system; then
			exit 1
		fi
		reboot_system "Firmware update finished"
	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
		;;
		--swap-active-system) SWAP_ACTIVE_SYSTEM=1
		;;
		-s) SCRIPT_MODE=1
		;;
		-v|--verbose) VERBOSE="-v"
		;;
		-h|--help) usage;exit
		;;
		*) UPDATE_FILE="${1}"
		    break
		;;
	esac
	shift
done

get_active_system || exit

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

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

# Swap active system.
if [ ${SWAP_ACTIVE_SYSTEM} -eq 1 ]; then
	if ! swap_active_system; then
		exit 1
	fi
	reboot_system "Swapped active system to $(echo ${ALT_BOOT} | cut -d'_' -f2 | tr [:lower:] [:upper:])"
	exit
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
