#!/bin/sh

# Copyright: (C) 2017-2021, Stefan Lippers-Hollmann <s.l-h@gmx.de>

# License: ISC
# Permission to use, copy, modify, and/or distribute this software for any
# purpose with or without fee is hereby granted, provided that the above
# copyright notice and this permission notice appear in all copies.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
# ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
# WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
# ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
# OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

set -e

usage() {
	cat <<EOF
$(basename $0): parse and set the dual boot configuration on the ZyXEL NBG6817

Parameters:
-h|--help               this usage information
-hh|--full-help         also display advanced usage informations
-l|--list               list installed firmwares
-sl|--short-list        abbreviated list installed firmwares
--toggle-rootfs         toggle the stored bootflag to boot from the alternative
                        partition set.
                        please reboot immediately after using this parameter!

EOF
}

advanced_usage() {
	usage
	cat <<EOF
Advanced parameters:
--get-rootfs            get the currently configured rootfs
                        /dev/mmcblk0p5 or /dev/mmcblk0p8
--set-rootfs <blkdev>   set the desired rootfs
                        /dev/mmcblk0p5 or /dev/mmcblk0p8
--fake-rootfs           change /proc/cmdline to display the alternative rootfs
                        instead of the currently booted one, this will instruct
                        sysupgrade to flash the currently booted partition set,
                        instead of the alternative one.
--unfake-rootfs         restore the original /proc/cmdline, to allow writing
                        to the alternative partition set again, if overridden
--check-mtd-integrity   check the integrity of the dualflag mtd
--get-mtd               get the dualflag mtd
                        /dev/mtdblock6 (ZyXEL OEM) or /dev/mtdblock11 (OpenWrt)
--get-bootloader-rootfs get the rootfs chosen by the bootloader from sysfs
                        /dev/mmcblk0p5 or /dev/mmcblk0p8
--get-cmdline-rootfs    get the rootfs configured by /proc/cmdline
                        /dev/mmcblk0p5 or /dev/mmcblk0p8
--get-version <blkdev>  try to decode the version of the firmware installed in
                        the named rootfs (/dev/mmcblk0p5 or /dev/mmcblk0p8)
--reset-rootfs <blkdev> DANGEROUS: rewrite the dualflag mtd completely and set
                        the desired rootfs
                        /dev/mmcblk0p5 or /dev/mmcblk0p8

EOF
}

get_mtd() {
	local DUAL_FLAG_MTD

	# OpenWrt with kernel >=5.10 uses "0:dual_flag" (/dev/mtdblock11)
	DUAL_FLAG_MTD="$(find_mtd_part 0:dual_flag)"
	if [ -z "${DUAL_FLAG_MTD}" ]; then
		# OpenWrt with kernel <5.10 uses "0:DUAL_FLAG" (/dev/mtdblock11)
		DUAL_FLAG_MTD="$(find_mtd_part 0:DUAL_FLAG)"
		if [ -z "${DUAL_FLAG_MTD}" ]; then
			# ZyXEL's OEM firmware uses "dualflag" (/dev/mtdblock6)
			DUAL_FLAG_MTD="$(find_mtd_part dualflag)"
		fi
	fi

	# check that the detected mtd is really a block device
	if [ -b "${DUAL_FLAG_MTD}" ]; then
		echo $DUAL_FLAG_MTD
	else
		echo "ERROR: failing to find DUAL_FLAG mtd" >&2
		exit 1
	fi
}

get_virtual_rootfs() {
	if [ -r "$1" ]; then
		for param in $(cat ${1}); do
			case "${param}" in
				root=*)
					echo "${param#root=}"
					;;
			esac
		done
	else
		echo "ERROR: failing to read ${1}." >&2
		exit 14
	fi
}

unfake_rootfs() {
	# if previously overmounted, umount /proc/cmdline
	while :; do umount /proc/cmdline 2>/dev/null || break ; done
}

fake_rootfs() {
	local FAKE_CMDLINE_FILE
	local FAKE_ROOTFS

	unfake_rootfs

	FAKE_CMDLINE_FILE="$(mktemp -qt cmdline.XXXXXX)"
	if [ ! -w "${FAKE_CMDLINE_FILE}" ]; then
		echo "ERROR: failing to create ${FAKE_CMDLINE_FILE}." >&2
		exit 15
	fi

	case "$(get_virtual_rootfs /sys/firmware/devicetree/base/chosen/bootloader-args)" in
		/dev/mmcblk0p5)
			FAKE_ROOTFS="mmcblk0p8"
			;;
		/dev/mmcblk0p8)
			FAKE_ROOTFS="mmcblk0p5"
			;;
	esac

	sed "s/root=[\/a-z0-9]*/root=\\/dev\\/${FAKE_ROOTFS}/" /proc/cmdline \
		>"${FAKE_CMDLINE_FILE}"
	chmod 0444 "${FAKE_CMDLINE_FILE}"

	if ! mount --bind "${FAKE_CMDLINE_FILE}" /proc/cmdline; then
		echo "ERROR: failing to over-mount /proc/cmdline." >&2
		exit 16
	fi
}

get_rootfs() {
	local BOOT_FLAG

	BOOT_FLAG=$(dd if=$(get_mtd) bs=1 count=1 2>/dev/null | hexdump -n 1 -e '1/1 "%02x\n"')
	if [ "0x${BOOT_FLAG}" = "0xff" ]; then
		# using header (mmcblk0p3)/ kernel (mmcblk0p4)/ rootfs (mmcblk0p5)
		echo "/dev/mmcblk0p5"
	elif [ "0x${BOOT_FLAG}" = "0x01" ]; then
		# using header_1 (mmcblk0p6)/ kernel_1 (mmcblk0p7)/ rootfs_1 (mmcblk0p8)
		echo "/dev/mmcblk0p8"
	else
		echo "ERROR: can't parse bootflag"
		exit 2
	fi
}

get_version() {
	local HEADER
	local KERNEL
	local ROOTFS
	local ALTBOOT
	local BOOTFLAG
	ROOTFS="$1"

	case "$ROOTFS" in
		/dev/mmcblk0p5)
			HEADER=/dev/mmcblk0p3
			KERNEL=/dev/mmcblk0p4
			ALTBOOT=""
			;;
		/dev/mmcblk0p8)
			HEADER=/dev/mmcblk0p6
			KERNEL=/dev/mmcblk0p7
			ALTBOOT="_1"
			;;
		*)
			echo "ERROR: invalid rootfs supplied." >&2
			exit 13
			;;
	esac

	FIRMWARE_VERSION="$(dd if=${HEADER} bs=1 skip=8 count=16 2>/dev/null)"
	KERNEL_VERSION="$(dd if=${KERNEL} bs=1 skip=32 count=128 2>/dev/null | sed -n 's/.*\(^.*Linux-[0-9\\.]*\).*/\1/p')"

	case "$KERNEL_VERSION" in
		ARM*Linux-*)
			# LEDE/ OpenWrt
			FIRMWARE_VERSION="LEDE/ OpenWrt <unknown revision>"
			;;
		Linux-*)
			# ZyXEL OEM firmware
			:
			;;
		*)
			echo "ERROR: version of the installed firmware can't be decoded." >&2
			exit 12
			;;
	esac

	if [ "$(get_rootfs)" = "${ROOTFS}" ]; then
		BOOTFLAG="yes"
	else
		BOOTFLAG="no"
	fi

	cat <<EOF
"header${ALTBOOT}" partition is located at "${HEADER}"
"kernel${ALTBOOT}" partition is located at "${KERNEL}"
"rootfs${ALTBOOT}" partition is located at "${ROOTFS}"

FIRMWARE_VERSION="${FIRMWARE_VERSION}"
KERNEL_VERSION="${KERNEL_VERSION}"
UNAME_VERSION="${KERNEL_VERSION##*Linux-}"
BOOTFLAG="${BOOTFLAG}"
EOF
}

get_short_list() {
	local PRIMARY_FIRMWARE_VERSION
	local PRIMARY_KERNEL_VERSION
	local SECONDARY_FIRMWARE_VERSION
	local SECONDARY_KERNEL_VERSION
	local BOOT_PARTITION_SET

	PRIMARY_FIRMWARE_VERSION="$(dd if=/dev/mmcblk0p3 bs=1 skip=8 count=16 2>/dev/null)"
	PRIMARY_KERNEL_VERSION="$(dd if=/dev/mmcblk0p4 bs=1 skip=32 count=128 2>/dev/null | sed -n 's/.*\(^.*Linux-[0-9\\.]*\).*/\1/p')"
	case "$PRIMARY_KERNEL_VERSION" in ARM*Linux-*) PRIMARY_FIRMWARE_VERSION="OpenWrt" ;; esac

	SECONDARY_FIRMWARE_VERSION="$(dd if=/dev/mmcblk0p6 bs=1 skip=8 count=16 2>/dev/null)"
	SECONDARY_KERNEL_VERSION="$(dd if=/dev/mmcblk0p7 bs=1 skip=32 count=128 2>/dev/null | sed -n 's/.*\(^.*Linux-[0-9\\.]*\).*/\1/p')"
	case "$SECONDARY_KERNEL_VERSION" in ARM*Linux-*) SECONDARY_FIRMWARE_VERSION="OpenWrt" ;; esac

	case "$(get_rootfs)" in
		/dev/mmcblk0p5)
			BOOT_PARTITION_SET="primary"
			;;
		/dev/mmcblk0p8)
			BOOT_PARTITION_SET="secondary"
			;;
	esac

	cat <<EOF
$(basename $0): short listing of the installed firmware sets

Primary partition set (/dev/mmcblk0p5):
	firmware version: ${PRIMARY_FIRMWARE_VERSION}
	kernel version:   ${PRIMARY_KERNEL_VERSION##*Linux-}

Secondary partition set (/dev/mmcblk0p8):
	firmware version: ${SECONDARY_FIRMWARE_VERSION}
	kernel version:   ${SECONDARY_KERNEL_VERSION##*Linux-}

Boot from:
	${BOOT_PARTITION_SET}

EOF

	check_for_fake_rootfs
}

check_for_fake_rootfs() {
	local ROOT_FROM_CMDLINE
	local ROOT_FROM_BOOTLOADER

	ROOT_FROM_CMDLINE="$(get_virtual_rootfs /proc/cmdline)"
	ROOT_FROM_BOOTLOADER="$(get_virtual_rootfs /sys/firmware/devicetree/base/chosen/bootloader-args)"
	if [ "${ROOT_FROM_CMDLINE}" != "${ROOT_FROM_BOOTLOADER}" ]; then
		cat <<EOF
WARNING: /proc/cmdline has been overriden, sysupgrade will flash the currently
         booted partition set instead of the alternative one.
         sysupgrade will overwrite: /dev/mmcblk0p$(( ${ROOT_FROM_CMDLINE#\/dev\/mmcblk0p} - 1)) and ${ROOT_FROM_CMDLINE}

The original behaviour of sysupgrade writing to the alternative partition set
can be restored via "$(basename $0) --unfake-rootfs".

EOF
	fi
}

set_rootfs() {
	local ROOTFS
	ROOTFS="$1"

	if [ -z "${ROOTFS}" ] || [ ! -b "${ROOTFS}" ]; then
		echo "ERROR: provided rootfs (${ROOTFS}) not a block device" >&2
		exit 3
	fi

	case "$ROOTFS" in
		/dev/mmcblk0p5)
			printf "\xff" >$(get_mtd)
			;;
		/dev/mmcblk0p8)
			printf "\x01" >$(get_mtd)
			;;
		*)
			echo "ERROR: invalid rootfs (${ROOTFS})" >&2
			exit 4
			;;
	esac
}

toggle_rootfs() {
	local ROOTFS

	ROOTFS="$(get_rootfs)"
	echo "Current rootfs: ${ROOTFS}"

	case "$ROOTFS" in
		/dev/mmcblk0p5)
			ROOTFS="/dev/mmcblk0p8"
			;;
		/dev/mmcblk0p8)
			ROOTFS="/dev/mmcblk0p5"
			;;
	esac

	set_rootfs "${ROOTFS}"
	echo "New rootfs: ${ROOTFS}"
	echo ""
	echo "Please reboot now."
}

reset_rootfs() {
	local ROOTFS
	ROOTFS="$1"

	if [ -z "${ROOTFS}" ] || [ ! -b "${ROOTFS}" ]; then
		echo "ERROR: provided rootfs (${ROOTFS}) not a block device" >&2
		exit 3
	fi

	case "$ROOTFS" in
		/dev/mmcblk0p5)
			for i in $(seq 0 65535); do
				printf "\xff"
			done >$(get_mtd)
			;;
		/dev/mmcblk0p8)
			for i in $(seq 0 65535); do
				[ "${i}" -eq 0 ] && printf "\x01" || printf "\xff"
			done >$(get_mtd)
			;;
		*)
			echo "ERROR: invalid rootfs (${ROOTFS})" >&2
			exit 4
			;;
	esac
}

check_mtd_integrity() {
	local MTD_CHECKSUM

	# use md5 as checksum algorithm, sha256 is not supported by the
	# ZyXEL OEM firmware.
	MTD_CHECKSUM="$(md5sum $(get_mtd) | awk '{print $1}')"

	case $MTD_CHECKSUM in
		ecb99e6ffea7be1e5419350f725da86b)
			echo "valid dualflag signature found (0xFF, /dev/mmcblk0p5)."
			;;
		e107d3d780e73f0b5c7d48ec749e66f9)
			echo "valid dualflag signature found (0x01, /dev/mmcblk0p8)."
			;;
		*)
			echo "ERROR: unrecognized dualflag signature."
			exit 5
			;;
	esac
}

# ugly workaround, the OEM firmware doesn't define this function, while still
# providing and populating /tmp/sysinfo/board_name correctly, let this be
# redefined by /lib/functions.sh on LEDE/ OpenWrt
board_name() {
	[ -e /tmp/sysinfo/board_name ] && cat /tmp/sysinfo/board_name || echo "generic"
}

# provide find_mtd_part(), available in both LEDE/ OpenWrt and the ZyXEL OEM firmware
if [ -r /lib/functions.sh ]; then
	. /lib/functions.sh
else
	echo "ERROR: this tool can only be used on OpenWrt or the ZyXEL NBG6817 OEM firmware." >&2

	echo " "
	usage

	exit 6
fi

# bail out screaming, if this gets not executed on a ZyXEL NBG6817
case "$(board_name)" in
	zyxel,nbg6817|nbg6817)
		:
		;;
	*)
		echo "ERROR: this tool is only safe to be used on the ZyXEL NBG6817" >&2
		exit 7
		;;
esac

# at least one parameter is required
if [ -z "${1}" ]; then
	usage
	exit 8
fi

case "$1" in
	-h|--help)
		usage
		exit 0
		;;
	-hh|--full-help)
		advanced_usage
		exit 0
		;;
	--check-mtd-integrity)
		echo $(check_mtd_integrity)
		exit 0
		;;
	--fake-rootfs)
		fake_rootfs
		exit 0
		;;
	--get-mtd)
		echo $(get_mtd)
		exit 0
		;;
	--get-bootloader-rootfs)
		echo $(get_virtual_rootfs /sys/firmware/devicetree/base/chosen/bootloader-args)
		exit 0
		;;
	--get-cmdline-rootfs)
		echo $(get_virtual_rootfs /proc/cmdline)
		exit 0
		;;
	--get-rootfs)
		echo $(get_rootfs)
		exit 0
		;;
	--get-version)
		if [ -n "$2" ]; then
			get_version "$2"
		else
			echo "ERROR: rootfs not provided" >&2
			usage
			exit 11
		fi
		exit 0
		;;
	-l|--list)
		get_version /dev/mmcblk0p5
		echo ""
		echo ""
		get_version /dev/mmcblk0p8
		echo ""
		echo ""
		check_for_fake_rootfs
		exit 0
		;;
	-sl|--short-list)
		get_short_list
		exit 0
		;;
	--set-rootfs)
		if [ -n "$2" ]; then
			set_rootfs "$2"
		else
			echo "ERROR: rootfs not provided" >&2
			usage
			exit 9
		fi
		exit 0
		;;
	--reset-rootfs)
		if [ -n "$2" ]; then
			reset_rootfs "$2"
		else
			echo "ERROR: rootfs not provided" >&2
			usage
			exit 9
		fi
		exit 0
		;;
	--toggle-rootfs)
		toggle_rootfs
		exit 0
		;;
	--unfake-rootfs)
		unfake_rootfs
		exit 0
		;;
	*)
		usage
		exit 10
		;;
esac