#!/bin/ash
# Simple boot menu selector using kexec
# (C) James Budiono 2013
# License: GNU GPL Version 3 or later
#
# Features: auto-boot, menu file editing, menu comments
#
# box-drawing idea from:
# http://top-scripts.blogspot.com/2011/01/power-of-echo-command-bash-console.html
#

### configuration
MENUFILE=bootmenu.cfg
EVFIFO=/tmp/evfifo.$$
APPINFO="BootMenu 1.0"
TMP_MOUNT=/tmp/bootmnt
EXCLUDED_DEVICES="^nbd|^ram|^nand|^loop" # devices we don't scan for bootmenu

MENU_TOP=3     # reserve top 3 lines, including top border 
MENU_BOTTOM=-9 # reserve bottom 9 lines, including bottom border
MENU_HEIGHT=$(($MENU_BOTTOM-$MENU_TOP-1)) # in lines, excluding top and bottom border
AUTOBOOT_LINE=$(($MENU_BOTTOM + 1))
APPEND_LINE=$(($MENU_BOTTOM + 2))  # root, kernel, initrd, append (two lines)
COMMENT_LINE=$(($MENU_BOTTOM + 7)) # comment

### run-time variables
NUMENTRIES=0
DEFAULT_ENTRY=1
BOOTDELAY=10 # auto-boot delay


#################### helpers ####################

### terminal controls
trap 'if [ $$ -ne 1 ]; then cleanup 2> /dev/null; exit; fi' INT QUIT
TAB=$(printf "\011")
ESC=$(printf "\033")
BOLD="$ESC[1m"
NORM="$ESC[0m"
TPUT()   { printf "${ESC}[${1};${2}f"; } # y,x. ESC [f better than ESC [H
CLEAR()  { printf "${ESC}c"; }
CIVIS()  { printf "${ESC}[?25l"; }
CNORM()  { printf "${ESC}[?12l${ESC}[?25h"; }
DRAW()   { [ $TERM = "linux" ] && printf "${ESC}%%@"; printf "${ESC}(0"; }
WRITE()  { printf "${ESC}(B"; }
KILL()   { printf "${ESC}[2K"; } # delete line
ERASE()  { printf "${ESC}[J"; }	 # delete end of screen

MARK()   { printf "${ESC}[7m"; }
UNMARK() { printf "${ESC}[27m"; }
BLUE()   { printf "${ESC}c${ESC}[H${ESC}[J${ESC}[37;44m${ESC}[J"; }


### drawing routines

# returns MAXX (horizontal length), MAXY (vertical length)
get_ttysize() {
	set -- $(ttysize)
	MAXX=$1
	MAXY=$2
}


### draw horizontal lines for menu border
# $1 - row location, MAXX
draw_horz_line() {
	local x=2 y=$1
	TPUT $y $x # y, x
	DRAW
	while [ $x -lt $MAXX ]; do
		printf "q"
		x=$(($x+1))
	done
	WRITE
}

### draw vertical lines for menu border
# $1 - starting row, $2 - length (in lines) to draw, MAXX
draw_vert_line() {
	local y=$1 len=$2
	DRAW
	while [ $len -gt 0 ]; do
		TPUT $y 1; printf "x"
		TPUT $y $MAXX; printf "x"
		y=$(($y+1)); len=$(($len-1))
	done
	WRITE
}

### draw the corners for menu borders
# MENU_TOP, MENU_BOTTOM, MAXX
draw_corners() {
	DRAW
	TPUT $MENU_TOP 1; printf "l"
	TPUT $MENU_TOP $MAXX; printf "k"
	TPUT $MENU_BOTTOM 1; printf "m"
	TPUT $MENU_BOTTOM $MAXX; printf "j"
	WRITE
}

### draw the menu border  in full
draw_menu_border() {
	draw_horz_line $MENU_TOP
	draw_horz_line $MENU_BOTTOM
	draw_vert_line $(($MENU_TOP+1)) $MENU_HEIGHT
	draw_corners $MENU_TOP $MENU_BOTTOM
}

### draw the menu entries with the highlight cursor
# $1 - offset, $2 - cursor location (must be between 1 to MENU_HEIGHT)
# MENU_HEIGHT, MENU_TOP, MAXX
draw_menu_entries() {
	local offset=$1 cursor=$2 entry=1
	local height=$MENU_HEIGHT top=$(($MENU_TOP+1))
	
	# clear the entries first
	while [ $height -gt 0 ]; do
		TPUT $top 2; KILL
		height=$((height-1))
		top=$(($top+1))	
	done
	height=$MENU_HEIGHT top=$(($MENU_TOP+1))
	
	# and print the new ones
	while read -r p; do
		case "$p" in
			title*) echo "${p#title}" ;;
		esac
	done < /tmp/$MENUFILE |
	while read -r p; do
		if [ $offset -ne 0 ]; then
			offset=$((offset-1))
			continue
		fi
		if [ $height -gt 0 ]; then
			TPUT $top 2
			if [ $entry -eq $cursor ]; then
				MARK
				printf "%*s" $(($MAXX-2)) " "
				TPUT $top 2
				printf "%s" "$p"
				UNMARK
			else
				printf "%s" "$p"			
			fi
			height=$((height-1))
			top=$(($top+1))
		else 
			break
		fi
		entry=$(($entry+1))
	done
	
	# reconstruct damaged borders
	draw_vert_line $(($MENU_TOP+1)) $MENU_HEIGHT	
}

###
# draw comments for a particular menu entry
# $1-entry
draw_menu_comment() {
	local p current=0
	TPUT $APPEND_LINE 1; ERASE
	while read -r p; do
		[ $current -gt $entry ] && return # optimise
		if [ $current -eq $entry ]; then
			case "$p" in
				title*) current=$(($current+1)) ;;
				root*)    TPUT $APPEND_LINE 1; echo "$p" ;;
				kernel*)  TPUT $(($APPEND_LINE+1)) 1; echo "$p" ;;
				initrd*)  TPUT $(($APPEND_LINE+2)) 1; echo "$p" ;;
				append*)  TPUT $(($APPEND_LINE+3)) 1; echo "$p" ;;
				comment*) TPUT $COMMENT_LINE 1; echo "${p#comment}" | sed 's/^[ \t]*//' ;;
				
				# special entries
				reboot*)   TPUT $APPEND_LINE 1; echo "Exit and reboot the system." ;;
				init*)     TPUT $APPEND_LINE 1; echo "Exit and run /init to start system with current parameters." ;;
				poweroff*) TPUT $APPEND_LINE 1; echo "Exit and power off the system." ;;
				shell*)    TPUT $APPEND_LINE 1; echo "Exit and run PID 1 shell." ;;
			esac
		else
			case "$p" in
				title*) current=$(($current+1)) ;;
			esac		
		fi
	done < /tmp/$MENUFILE
}

### print centered text
# $1 - row, $2 - text, MAXX
print_center() {
	local len=${#2}
	local xstart=$(( ($MAXX-$len)/2 ))
	TPUT $1 $xstart
	printf "$2"
}

################ event handling ##############
### 
# get key pressed event
# returns "." or key code 
get_key() {	
	# get key code or returns dot if there is no input
	local p=.
	IFS="" read -s -n 1 -t 1 p # one second wait
	case "$p" in
		.) echo . ;;
		I|i) echo I ;;
		E|e) echo E ;;
		R|r) echo R ;;
		P|p) echo P ;;
		"$TAB") echo TAB ;;
		"") echo ENTER ;;		
		"$ESC") p=.
				read -n 1 -t 1 p
				case "$p" in
					"[") p=.
						 read -n 1 -t 1 p
						 case "$p" in
							A) echo UP ;;
							B) echo DOWN ;;
							C) echo RIGHT ;;
							D) echo LEFT ;;
							*) echo . ;;
						 esac ;;
					.) echo ESC ;;
					*) echo . ;;
				esac ;;
	esac
}


################# setup and cleanup ####################

######
# initial setup - mount proc & sys, find bootmenu.cfg
setup() {
	# mount -t proc proc /proc - so that busybox exec applet works
	if [ $$ -eq 1 ]; then
		/bin/mount -t proc proc /proc
		mount -t sysfs sysfs /sys # need to search bootmenu
		if ! mount -t devtmpfs devtmpfs /dev 2> /null; then		# /dev/null may not exist yet
			# if no devtmpfs, use tmpfs and use mdev instead
			mount -t tmpfs tmpfs /dev
			mdev -s
		fi		
	fi
	
	# honour waitdev= variable
	! [ -z "$waitdev" ] && echo "Waiting for devices ($waitdev seconds) ... " && 
	sleep $waitdev
	
	# find bootmenu.cfg
	mkdir $TMP_MOUNT 2> /dev/null
	rm /tmp/$MENUFILE 2> /dev/null
	printf "Looking for $MENUFILE in "	
	if [ -z "$bootmenu" ]; then
		set -- $(ls /sys/class/block | grep -Ev $EXCLUDED_DEVICES)
		while [ "$1" ]; do
			printf "$1 "
			if mount -o ro /dev/$1 $TMP_MOUNT 2>/dev/null; then
				[ -e $TMP_MOUNT/$MENUFILE ] && bootmenu=$1 && break
				umount /dev/$1 2>/dev/null
			fi
			shift
		done
	else
		printf "$bootmenu "
		mount -o ro /dev/$bootmenu $TMP_MOUNT 2>/dev/null
	fi
	[ -e $TMP_MOUNT/$MENUFILE ] && cp $TMP_MOUNT/$MENUFILE /tmp && printf " found.\n" ||
	printf "not found.\n"
	umount /dev/$bootmenu 2>/dev/null	

	# one-time visual setup
	get_ttysize
	MENU_BOTTOM=$(($MAXY+$MENU_BOTTOM))
	MENU_HEIGHT=$(($MAXY+$MENU_HEIGHT))
	AUTOBOOT_LINE=$(($MAXY+$AUTOBOOT_LINE))
	APPEND_LINE=$(($MAXY+$APPEND_LINE))
	COMMENT_LINE=$(($MAXY+$COMMENT_LINE))
}

### cleanup - clear temp files, kill helper apps, reset screen
cleanup() {
	TPUT $AUTOBOOT_LINE 1; ERASE
	rm /tmp/$MENUFILE
	stty echo
	CNORM
	if [ $$ -eq 1 ]; then
		umount /sys
		umount /proc
		umount /dev
	fi
}

### 
# visual_setup - get screen size, hid cursor, clear screen, print borders
# logo, etc
visual_setup() {
	CLEAR
	CIVIS
	stty -echo
	print_center 1 "$APPINFO"
	print_center 2 "Up/Down-Select Tab/E-Edit Enter-Boot Esc-Shell P-Poweroff I-Init R-Reboot"
	draw_menu_border	
}

### 
# read (and validate) boot menu file
load_boot_menu() {
	NUMENTRIES=0
	! [ -r /tmp/$MENUFILE ] && return
	while read -r p; do
		case "$p" in
			title*) NUMENTRIES=$(($NUMENTRIES+1)) ;;
			bootdelay*) [ $BOOTDELAY -ne 0 ] && BOOTDELAY=${p#bootdelay} ;;
			default*) DEFAULT_ENTRY=${p#default}
		esac
	done < /tmp/$MENUFILE
	[ $BOOTDELAY -ne 0 ] && BOOTDELAY=$((BOOTDELAY+1)) # quirk
}

###
# preload kernel for kexec
# $1-entry
# return: exitcode (if blank, failed)
preload_kernel() {
	local root= kernel= initrd= append= p current=0 dtb=
	
	# collect parameters
	exitcode=
	while read -r p; do
		[ $current -gt $entry ] && break
		if [ $current -eq $entry ]; then
			case "$p" in
				title*) current=$(($current+1)) ;;
				root*) root=$(echo ${p#root} | sed 's/[ \t]*//') ;;
				kernel*) kernel=$(echo ${p#kernel} | sed 's/[ \t]*//') ;;
				initrd*) initrd=$(echo ${p#initrd} | sed 's/[ \t]*//') ;;
				append*) append=$(echo ${p#append} | sed 's/[ \t]*//') ;;
				dtb*) append=$(echo ${p#dtb} | sed 's/[ \t]*//') ;;
				
				# special return codes
				shell*) exitcode=shell; return ;;
				init*) exitcode=init; return ;;
				poweroff*) exitcode=poweroff; return ;;
				reboot*) exitcode=reboot; return ;;
			esac
		else
			case "$p" in
				title*) current=$(($current+1)) ;;
			esac		
		fi
	done < /tmp/$MENUFILE
	
	# validate - must have root and kernel at least
	[ $root ] && [ $kernel ] || return 1 # fail
	
	# attempt to mount $root and obtain $kernel and $initrd
	mkdir $TMP_MOUNT 2> /dev/null
	! mount /dev/$root $TMP_MOUNT && return 1 # fail
	
	# check that files exist
	[ -e $TMP_MOUNT/$kernel ] && kernel="-l $TMP_MOUNT/$kernel" || kernel=
	[ $initrd ] && [ -e $TMP_MOUNT/$initrd ] && initrd="--initrd=$TMP_MOUNT/$initrd" || initrd=
	[ $dtb ] && [ -e $TMP_MOUNT/$dtb ] && dtb="--dtb=$TMP_MOUNT/$dtb" || dtb=--atags
		
	# pre-load if kernel exist
	[ "$kernel" ] && kexec $initrd --command-line="$append" $dtb $kernel && exitcode=boot
	umount $TMP_MOUNT
}

##################### the interactive main loop ##########################
cancel_autoboot() {
	BOOTDELAY=0
	TPUT $(($MENU_BOTTOM+1)) 1
	printf "%*s" "$MAXX" " "
}
#$1-default entry
main_loop() {
	local finish= offset=$((${1:-1}-1)) cursor=1 entry=1 preventry=0 autoboot=""
	
	while [ -z $finish ]; do
		entry=$(($offset + $cursor)) # 1-based
		if [ $entry -ne $preventry ]; then
			draw_menu_entries $offset $cursor
			draw_menu_comment
			preventry=$entry
		fi

		if [ -z "$autoboot" ]; then
			set -- $(get_key)
		else
			set -- ENTER
		fi
		case $1 in
			.)	# autoboot
				if [ $BOOTDELAY -ne 0 ]; then
					BOOTDELAY=$(($BOOTDELAY - 1))
					if [ $BOOTDELAY -eq 0 ]; then
						autoboot=true
						continue
					else
						TPUT $AUTOBOOT_LINE 1
						printf "Auto-boot $entry in $BOOTDELAY seconds ..."
					fi
				fi
				;;
			
			DOWN)
				# move selection down
				cancel_autoboot
				if [ $(($cursor + $offset)) -lt $NUMENTRIES ]; then
					cursor=$(($cursor+1))
					if [ $cursor -gt $MENU_HEIGHT ]; then
						cursor=$(($cursor-1))
						offset=$(($offset+1))
					fi
				fi
				;;
				
			UP)
				# move selection down
				cancel_autoboot
				cursor=$(($cursor-1))
				if [ $cursor -le 0 ]; then
					cursor=1
					offset=$(($offset-1))
					[ $offset -le 0 ] && offset=0
				fi
				;;
				
			ENTER)
				# boot
				cancel_autoboot
				TPUT $AUTOBOOT_LINE 1; KILL
				printf "Booting entry $entry ..."
				preload_kernel $entry 
				[ "$exitcode" ] && break
				TPUT $AUTOBOOT_LINE 1; KILL
				printf "Can't boot: unable to load kernel for selected entry ..."
				sleep 2
				preventry=0
				;;
				
			TAB|E)
				# edit menufile
				cancel_autoboot
				CNORM
				vi /tmp/$MENUFILE
				CIVIS
				exitcode=reload_menu
				break
				;;
				
			ESC)
				exitcode=shell
				break
				;;				
			P)
				exitcode=poweroff
				break
				;;
			I)
				exitcode=init
				break
				;;
			R)
				exitcode=reboot
				break
				;;
		esac
	done
}

################## main: program entry point ################
setup
#cp /root/bootmenu.cfg /tmp #stub
while true; do	
	load_boot_menu
	
	# see if bootmenu is any good
	if [ $NUMENTRIES -eq 0 ]; then
		if [ $$ -eq 1 ]; then
			cleanup 2> /dev/null					
			printf "$MENUFILE is bad or not found. Starting /init in 5 seconds ...\n"
			sleep 5
			exec /init
		else 
			printf "$MENUFILE is bad or not found, exiting.\n"
			exit
		fi
	fi
	
	# if yes, then let user choose
	visual_setup
	main_loop $DEFAULT_ENTRY
	
	# see what to do after this
	case $exitcode in
		reload_menu) 
			continue ;;
			
		boot)
			cleanup
			echo Booting new kernel ...
			kexec -f -e
			;;
		
		shell) 
			cleanup
			echo "Exiting to shell (PID 1)"
			exec /bin/ash
			;;
			
		poweroff)
			cleanup
			echo Powering off ...
			poweroff -f
			;;
			
		reboot)
			cleanup
			echo Rebooting ...
			reboot -f
			;;
		
		init)
			cleanup
			echo "Starting /init"
			exec /init
			;;
	esac
done