#!/bin/bash

function tftp_err()
{
	if [ "$quiet" = "0" ]; then
		echo "$1"
		exit 1
	fi
}

function tftp_fetch()
{
	if [ "$1" = "-q" ]; then
		quiet=1
		shift
	else
		quiet=0
	fi

	if [ ! -z "$local_dir" ]; then
		if [ -e $local_dir/$2 ]; then
			cp $local_dir/$2 $3
			return 0
		elif [ -e $local_dir/${2##*/} ]; then
			cp $local_dir/${2##*/} $3
			return 0
		else
			tftp_err "ERROR: can't find ${2##*/} in $local_dir"
			return 1
		fi
	fi
	echo "Fetching via tftp: $1:$2"
	if [ $tftp_block_size = 0 ]; then
		if ! tftp -g $1 -r $2 -l $3; then
			tftp_err "ERROR: tftp failed.  Aborting."
			return 1
	   fi
	else
		if ! tftp -g $1 -r $2 -l $3 -b $tftp_block_size; then
		   tftp_err "ERROR: tftp failed. Aborting."
		   return 1
	   fi
	fi
	return 0
}

function find_mtd_partition()
{
	arg=$1
	out=$2
	for x in /sys/class/mtd/mtd*; do
		if [ -e $x/name ]; then
			name=$(cat $x/name)

			# 0x400 = MTD_WRITEABLE
			flags=$(cat $x/flags)
			if [ "$name" = "$arg" \
					-a $(($flags & 0x400)) != 0 ]; then
				eval $out=${x#*mtd/mtd}
				return 0
			fi
		fi
	done
	return 1
}

function get_mtd_attr()
{
	arg0=$1
	arg1=$2
	out=$3
	file=/sys/class/mtd/mtd${arg0}/$arg1
	if [ -e $file ]; then
		retval=$(cat $file)
		eval $out=$retval
	else
		#echo $file
		return 1
	fi
	return 0
}

function flash_is_emmc()
{
    if [ -e /dev/mmcblk0boot0 ]; then
        flash_type="disk"
        return 1 # note the backwards shell return-code semantics
    fi
	return 0
}

function flash_is_nand()
{
	mtdX=$1
    if ! flash_is_emmc; then
         return 1 # note the backwards shell return-code semantics
    fi
	! get_mtd_attr $mtdX type flash_type && echo "Cannot get flash type" && exit 1
	if [ "$flash_type" = "nand" ] || [ "$flash_type" = "mlc-nand" ]; then
		return 0 # note the backwards shell return-code semantics
	else
		return 1
	fi
}

if [ -d "/proc/mtd" ]
then
	RG_BLOCK0=$(grep "rg0" /proc/mtd | cut -d':' -f 1)
	k=$(grep "kernel0" /proc/mtd | cut -d '"' -f 2)
	KERNEL_ROOT=${k%?}
	k=$(grep "devtree0" /proc/mtd | cut -d '"' -f 2)
	DTB_ROOT=${k%?}
	k=$(grep "rg0" /proc/mtd | cut -d '"' -f 2)
	RG_ROOT=${k%?}
	RG_NONVOL0_MTD_NUM=$(grep "rgnonvol0" /proc/mtd | cut -d':' -f 1)
fi

function get_startup_env_param()
{
    # find rg nonvol
    get_mtd_attr ${RG_NONVOL0_MTD_NUM:3} type nonvol_flash_type
    if [ "$nonvol_flash_type" = "nand" ] || [ "$nonvol_flash_type" = "mlc-nand" ]
    then #if nonvol is in nand, then board should have primary flash nand (nand boot)
        rgnonvol_type=ubifs
    elif [ "$nonvol_flash_type" = "nor" ]
    then #nor/nand or nor/emmc
        rgnonvol_type=jffs2
    else #not an mtd device, (emmc boot)
        rgnonvol_type=ext4
    fi

    # find rg apps
    get_mtd_attr ${RG_BLOCK0:3} type rg_flash_type
    if [ "$rg_flash_type" = "nand" ] || [ "$rg_flash_type" = "mlc-nand" ]
    then
        if [ "$UBIFS" -eq 1 ]; then
            root_path="root=ubi0:rootfs"
            rootfs_type=ubifs
            apps_type=ubifs
        else
            ubiblock="ubi.block=0,rootfs"
            root_path="root=/dev/ubiblock0_0"
            rootfs_type=squashfs
        fi
    else
        if [ "$nonvol_flash_type" = "nor" ]
        then #nor/emmc
            root_path="root=/dev/mmcblk0p5"
        else #emmc only
            root_path="root=/dev/mmcblk0p14"
        fi
        rootfs_type=ext4
        apps_type=ext4
    fi
}

function check_root()
{
	# make sure that we are not reformatting the active rootfs
	if [ -h /dev/root ]; then
		rootdev=$(readlink /dev/root)
		while [ ! -z $1 ]; do
			if [[ $rootdev = ${1}* ]]; then
				echo ""
				echo "ERROR: please run stbutil from the initrd kernel."
				echo ""
				exit 1
			fi
			shift
		done
	fi
}

get_magic_word() {
    dd if=$1 skip=13 bs=1 count=1 2>/dev/null | hexdump -v -n 1 -e '1/1 "%02x"'
}

is_squashfs() {
    if [ "$( get_magic_word $1 )" = "01" ]; then
        echo "1"
    fi
}

function handle_opt_1()
{
    if tftp_fetch  $tftphost $tftppath /tmp/rootfs.img; then
        if [ "$( is_squashfs /tmp/rootfs.img )" = "1" ]; then
            echo "Monolith header indicates UBI/squashfs file system"
            UBIFS=0
        else
            echo "Monolith header indicates UBIFS file system"
            UBIFS=1
        fi
		if [ "7145" == "${plat:0:4}" ] || [ "3390" == "${PLAT:0:4}" ]; then
			/bin/rprogramstore -f /tmp/rootfs.img
		elif [ "arm64" == "${PLAT:0:5}" ]; then
			/bin/rdload -f /tmp/rootfs.img
		else
			echo "Monolith not supported in $PLAT"
		fi
    fi

    echo "Finished."
    echo ""
    exit 0
}

function handle_opt_2()
{
	if [ "3390" == "${PLAT:0:4}" ]; then
		snmpset -v2c -cprivate 172.31.255.45 1.3.6.1.4.1.4413.2.99.1.1.2.4.2.2.2.2.0 a $tftphost
		echo "snmpset -v2c -cprivate 172.31.255.45 1.3.6.1.4.1.4413.2.99.1.1.2.4.2.2.2.2.0 a $tftphost"
		snmpset -v2c -cprivate 172.31.255.45 1.3.6.1.4.1.4413.2.99.1.1.2.4.2.2.2.3.0 s $tftppath
		echo "snmpset -v2c -cprivate 172.31.255.45 1.3.6.1.4.1.4413.2.99.1.1.2.4.2.2.2.3.0 s $tftppath"
		snmpset -v2c -cprivate 172.31.255.45 1.3.6.1.4.1.4413.2.99.1.1.2.4.2.2.2.6.0 i 1
		echo "snmpset -v2c -cprivate 172.31.255.45 1.3.6.1.4.1.4413.2.99.1.1.2.4.2.2.2.6.0 i 1"

		echo "Finished."
		echo ""
	elif [ "arm64" == "${PLAT:0:5}" ]; then
		if tftp_fetch  $tftphost $tftppath /tmp/opt2.img; then
			dd if=/tmp/opt2.img of=/dev/flash-boltb conv=fsync iflag=fullblock oflag=direct bs=1M 2> /dev/null;
			echo lunswap boltb > /proc/driver/vfbio/cmd;
			rm /tmp/opt2.img
			echo "$tftppath successfully uploaded in backup partition and swaped the lun"
		fi
	fi
	exit 0
}

function handle_opt_3()
{
	if tftp_fetch  $tftphost $tftppath /tmp/opt3.img; then
		dd if=/tmp/opt3.img of=/dev/flash-devtreeb conv=fsync iflag=fullblock oflag=direct bs=1M 2> /dev/null;
		echo lunswap devtreeb > /proc/driver/vfbio/cmd;
		rm /tmp/opt3.img
		echo "$tftppath successfully uploaded in backup partition and swaped the lun"
	fi

	exit 0
}

function handle_opt_4()
{
	if tftp_fetch  $tftphost $tftppath /tmp/opt4.img; then
		dd if=/tmp/opt4.img of=/dev/flash-kernelb conv=fsync iflag=fullblock oflag=direct bs=1M 2> /dev/null;
		echo lunswap kernelb > /proc/driver/vfbio/cmd;
		rm /tmp/opt4.img
		echo "$tftppath successfully uploaded in backup partition and swaped the lun"
	fi

	exit 0
}

function handle_opt_5()
{
	if tftp_fetch  $tftphost $tftppath /tmp/opt5.img; then
		dd if=/tmp/opt5.img of=/dev/flash-rgb conv=fsync iflag=fullblock oflag=direct bs=1M 2> /dev/null;
		echo lunswap rgb > /proc/driver/vfbio/cmd;
		rm /tmp/opt5.img
		echo "$tftppath successfully uploaded in backup partition and swaped the lun"
	fi

	exit 0
}

function handle_opt_6()
{
	if tftp_fetch  $tftphost $tftppath /tmp/opt6.img; then
		dd if=/tmp/opt6.img of=/dev/flash-cmb conv=fsync iflag=fullblock oflag=direct bs=1M 2> /dev/null;
		echo lunswap cmb > /proc/driver/vfbio/cmd;
		rm /tmp/opt6.img
		echo "$tftppath successfully uploaded in backup partition and swaped the lun"
	fi

	exit 0
}

function handle_opt_7()
{
	if tftp_fetch  $tftphost $tftppath /tmp/opt7.img; then
		dd if=/tmp/opt7.img of=/dev/flash-cmblb conv=fsync iflag=fullblock oflag=direct bs=1M 2> /dev/null;
		echo lunswap cmblb > /proc/driver/vfbio/cmd;
		rm /tmp/opt7.img
		echo "$tftppath successfully uploaded in backup partition and swaped the lun"
	fi

	exit 0
}

function handle_opt_8()
{
	if tftp_fetch  $tftphost $tftppath /tmp/opt8.img; then
		dd if=/tmp/opt8.img of=/dev/flash-smcblb conv=fsync iflag=fullblock oflag=direct bs=1M 2> /dev/null;
		echo lunswap smcblb > /proc/driver/vfbio/cmd;
		rm /tmp/opt8.img
		echo "$tftppath successfully uploaded in backup partition and swaped the lun"
	fi

	exit 0
}

function handle_opt_9()
{
	if tftp_fetch  $tftphost $tftppath /tmp/opt9.img; then
		dd if=/tmp/opt9.img of=/dev/flash-smcosb conv=fsync iflag=fullblock oflag=direct bs=1M 2> /dev/null;
		echo lunswap smcosb > /proc/driver/vfbio/cmd;
		rm /tmp/opt9.img
		echo "$tftppath successfully uploaded in backup partition and swaped the lun"
	fi

	exit 0
}

function handle_opt_10()
{
	if tftp_fetch  $tftphost $tftppath /tmp/opt10.img; then
		dd if=/tmp/opt10.img of=/dev/flash-gfapb conv=fsync iflag=fullblock oflag=direct bs=1M 2> /dev/null;
		echo lunswap gfapb > /proc/driver/vfbio/cmd;
		rm /tmp/opt10.img
		echo "$tftppath successfully uploaded in backup partition and swaped the lun"
	fi

	exit 0
}

function handle_opt_11()
{
	if tftp_fetch  $tftphost $tftppath /tmp/opt11.img; then
		dd if=/tmp/opt11.img of=/dev/flash-amsb conv=fsync iflag=fullblock oflag=direct bs=1M 2> /dev/null;
		echo lunswap amsb > /proc/driver/vfbio/cmd;
		rm /tmp/opt11.img
		echo "$tftppath successfully uploaded in backup partition and swaped the lun"
	fi

	exit 0
}

function handle_opt_12()
{
	if tftp_fetch  $tftphost $tftppath /tmp/opt12.img; then
		dd if=/tmp/opt12.img of=/dev/flash-atab conv=fsync iflag=fullblock oflag=direct bs=1M 2> /dev/null;
		echo lunswap atab > /proc/driver/vfbio/cmd;
		rm /tmp/opt12.img
		echo "$tftppath successfully uploaded in backup partition and swaped the lun"
	fi

	exit 0
}

function handle_opt_13()
{
	if tftp_fetch  $tftphost $tftppath /tmp/opt13.img; then
		dd if=/tmp/opt13.img of=/dev/flash-asfb conv=fsync iflag=fullblock oflag=direct bs=1M 2> /dev/null;
		echo lunswap asfb > /proc/driver/vfbio/cmd;
		rm /tmp/opt13.img
		echo "$tftppath successfully uploaded in backup partition and swaped the lun"
	fi

	exit 0
}

function handle_opt_14()
{
	if tftp_fetch  $tftphost $tftppath /tmp/opt14.img; then
		dd if=/tmp/opt14.img of=/dev/flash-stab conv=fsync iflag=fullblock oflag=direct bs=1M 2> /dev/null;
		echo lunswap stab > /proc/driver/vfbio/cmd;
		rm /tmp/opt14.img
		echo "$tftppath successfully uploaded in backup partition and swaped the lun"
	fi

	exit 0
}

function handle_opt_0()
{
	bootstring=`boltenv -g STARTUP`
	bak_bootstring="$bootstring"

	if [ $? -eq 0 ]; then
		echo "FYI, your STARTUP variable is set to the following:"
		echo "$bootstring"
	fi

	read -p "Would you like to set STARTUP as default value? (y/n) " -n 1 -r

	if [[ $REPLY =~ ^[Yy]$ ]]
	then
		get_startup_env_param
		if flash_is_nand ${RG_BLOCK0:3}; then
			if [ "$UBIFS" -eq 1 ]; then
				#NAND only or NOR/NAND
				bootstring="load -nz -raw -addr=\$DT_ADDRESS -max=0x10000 $DTB_ROOT\$DT_IDX;\$PREBOOT;boot $KERNEL_ROOT\$KL_IDX \"ubi.mtd=$RG_ROOT\$RG_IDX rootfstype=$rootfs_type ${root_path} platformboot ${apps_type}_apps ${rgnonvol_type}_data coherent_pool=1M \$XARGS\""
			else
				bootstring="load -nz -raw -addr=\$DT_ADDRESS -max=0x10000 $DTB_ROOT\$DT_IDX;\$PREBOOT;boot $KERNEL_ROOT\$KL_IDX \"ubi.mtd=$RG_ROOT\$RG_IDX $ubiblock rootfstype=$rootfs_type ${root_path} ro platformboot ${rgnonvol_type}_data coherent_pool=1M \$XARGS\""
			fi
		else
			if [ $nonvol_flash_type = "nor" ]; then
				DTB_ROOT="emmcflash.devtree"
				KERNEL_ROOT="emmcflash0.kernel"
			else
				DTB_ROOT="flash0.devtree"
				KERNEL_ROOT="flash0.kernel"
			fi
			#eMMC only or NOR/eMMC
			bootstring="load -nz -raw -addr=\$DT_ADDRESS -max=0x10000 $DTB_ROOT\$DT_IDX;\$PREBOOT;boot $KERNEL_ROOT\$KL_IDX \"${root_path} rootfstype=$rootfs_type mmc_core.removable=0 rootwait platformboot ${apps_type}_apps ${rgnonvol_type}_data coherent_pool=1M \$XARGS\""
		fi
		echo
		if [ "$bak_bootstring" != "$bootstring" ]; then
			echo "Setting BAK_STARTUP to \"$bak_bootstring\""
			boltenv -s BAK_STARTUP "$bak_bootstring"
		else
			echo "Not overwriting BAK_STARTUP with duplicate value"
		fi
		echo
		echo "Setting STARTUP to \"$bootstring\""
		boltenv -s STARTUP "$bootstring"
	fi
	exit 0
}

function usage()
{
	cat << EOF

dload v$vers - Platform software download utility
Copyright (C) 2016-2019 Broadcom Inc.

Usage: dload [ -a <selection> ] [ <tftphost>[:<dir>] ]

Options:
	-a <selection>		Use SELECTION as the menu item selection,
				instead of prompting the user
	-u			Use UBIFS settings
	<tftphost>		TFTP server hostname
	<dir>			Directory (relative path) on TFTP server

Examples:
	dload -a1 10.24.4.254:tftproot/monolith.bin

EOF
	exit 1
}

#
# MAIN
#

vers=1.2
rootfs_part=
rootfs_mtd_name=
rg_part=
rg_mtd_name=
stb_part=
stb_mtd_name=
kernel_part=
flash_type=
flash_erasesize=
flash_writesize=
hdd_dev=/dev/sda
hyper=0
: ${UBIFS:=0}

local_dir=
resp=
auto_resp=0
tftp_block_size=0

if [ ! -e /etc/brcmstb.conf ]; then
	echo "ERROR: missing /etc/brcmstb.conf"
else
	source /etc/brcmstb.conf
fi

#Source rcS.util
. /etc/init.d/rcS.util

# use default settings from brcmstb.conf
tftphost=$TFTPHOST
tftppath=$TFTPPATH
plat=$PLAT

if [ "7145" == "${plat:0:4}" ] || [ "3390" == "${PLAT:0:4}" ]
then
	for part_name in rootfs flash{2,1,0}.cm{,0,1}; do
		if find_mtd_partition $part_name rootfs_part; then
			if \
				! get_mtd_attr $rootfs_part type flash_type || \
				! get_mtd_attr $rootfs_part erasesize flash_erasesize || \
				! get_mtd_attr $rootfs_part writesize flash_writesize; then

				echo "Can't get flash info for rootfs"
				exit 1
			elif [ "$flash_type" = "ubi" ]; then
				# Don't use a UBI partition; must be the bare MTD
				continue
			else
				rootfs_mtd_name=$(cat /sys/class/mtd/mtd${rootfs_part}/name)
				break
			fi
		fi
	done
fi

if [ "7145" == "${plat:0:4}" ] || [ "3390" == "${PLAT:0:4}" ]
then
	for part_name in rootfs flash{2,1,0}.rg{,0,1}; do
		if find_mtd_partition $part_name rg_part; then
			if \
				! get_mtd_attr $rg_part type flash_type || \
				! get_mtd_attr $rg_part erasesize flash_erasesize || \
				! get_mtd_attr $rg_part writesize flash_writesize; then

				echo "Can't get flash info for rg"
				exit 1
			elif [ "$flash_type" = "ubi" ]; then
				# Don't use a UBI partition; must be the bare MTD
				continue
			else
				rg_mtd_name=$(cat /sys/class/mtd/mtd${rg_part}/name)
				break
			fi
		fi
	done
fi

if [ "3390" != "${PLAT:0:4}" ]
then
	for part_name in rootfs flash{2,1,0}.stb{,0,1}; do
		if find_mtd_partition $part_name stb_part; then
			if \
				! get_mtd_attr $stb_part type flash_type || \
				! get_mtd_attr $stb_part erasesize flash_erasesize || \
				! get_mtd_attr $stb_part writesize flash_writesize; then

				echo "Can't get flash info for stb"
				exit 1
			elif [ "$flash_type" = "ubi" ]; then
				# Don't use a UBI partition; must be the bare MTD
				continue
			else
				stb_mtd_name=$(cat /sys/class/mtd/mtd${stb_part}/name)
				break
			fi
		fi
	done
fi

if [ "7145" == "${plat:0:4}" ] || [ "3390" == "${PLAT:0:4}" ]
then
	for part_name in kernel flash{2,1,0}.kernel{,0,1}; do
		if find_mtd_partition $part_name kernel_part; then
			kernel_part_name=${part_name%?}
			break
		fi
	done
fi

if [ "7145" == "${plat:0:4}" ] || [ "3390" == "${PLAT:0:4}" ]
then
#devtree partition
	for part_name in kernel flash{2,1,0}.devtree{,0,1}; do
		if find_mtd_partition $part_name devtree_part; then
			dt_part_name=${part_name%?}
			break
		fi
	done
fi

# parse command line

while getopts "a:b:hu" opt; do
	case $opt in
		a)
			resp=$OPTARG;
			auto_resp=1
			;;
		b)
			tftp_block_size=$OPTARG
			;;
		u)
			UBIFS=1
			resp=$OPTARG;
			auto_resp=1
			;;
		h|?)
			usage
			;;
	esac
done
shift $(($OPTIND - 1))
if [ ! -z "$1" ]; then
	tftparg="$1"
	if [[ "$tftparg" = *:* ]]; then
		tftphost=${tftparg%%:*}
		tftppath=${tftparg#*:}
	else
		tftphost=$tftparg
		tftppath=""
	fi
fi

# construct menu

while :; do

	echo ""
	echo "dload v$vers"
	echo "------------"
	echo ""
	if [ -z "$local_dir" ]; then
		echo "Using TFTP server:     $tftphost"
		echo "Using TFTP path:       $tftppath"
	fi
	echo "Linux build target:    $plat"

	echo ""
	echo "Primary Linux flash:   ${flash_type:-none}"

	echo "UBIFS mode:     $UBIFS"
	echo ""

	if [ "3390" == "${PLAT:0:4}" ]
	then
		echo "1) Install monolith image via Ethernet"
		echo "2) Install monolith image via factory MIBs"
	fi

	if [ "arm64" == "${PLAT:0:5}" ]
	then
		echo "1) Install monolith image via Ethernet"
		echo "2) Install bolt image"
		echo "3) Install Device tree image"
		echo "4) Install kernel image"
		echo "5) Install RG image"
		echo "6) Install CM image"
		echo "7) Install CM bolt image"
		echo "8) Install SMCbootl image"
		echo "9) Install SMCOS image"
		echo "10) Install GFAP image"
		echo "11) Install AMS(SMON) image"
		echo "12) Install ATA image"
		echo "13) Install ASF image"
		echo "14) Install STA image"
	fi

	echo "q) Exit"
	echo ""
	echo -n "Selection: "

	if [ -z "$resp" ]; then
		read resp
	else
		echo $resp
	fi
	echo ""

	if [ "7145" == "${PLAT:0:4}" ] || [ "3390" == "${PLAT:0:4}" ]
	then
		start_arg=0
		max_arg=2
	else
		start_arg=1
		max_arg=7
	fi

	max_arg_tens=$((max_arg / 10))
	max_arg_ones=$((max_arg % 10))
	if [ $max_arg_tens -ne 0 ]
	then
		max_arg_single=9
	else
		max_arg_single=$max_arg
	fi

	case $resp in
		[$start_arg-$max_arg_single]|[1-$max_arg_tens][0-$max_arg_ones])
			sufname='$'suf${resp}
			eval suf=$sufname
			if [[ "$suf" = *"not available"* ]]; then
				echo "ERROR: option is not available on this system"
			else
				eval handle_opt_${resp}
			fi
			;;
		q|0)
			if [ "7145" == "${PLAT:0:4}" ] || [ "3390" == "${PLAT:0:4}" ]
			then
				handle_opt_0
			fi
			exit 0
			;;
		2|3|4|5|6|7|8|9|10|11|12|13|14)
			if [ "arm64" == "${PLAT:0:5}" ]
			then
				handle_opt_${resp};
			fi
			exit 0
			;;
		*)
			echo "ERROR: invalid selection"
			;;
	esac
	resp=""
	if [ $auto_resp = 1 ]; then
		exit 0
	fi
done
