#!/bin/bash

skip_modules=0
cmdline_file=""
target_host=""
img_file=""
def_cmdline=""
partition=4

function usage
{
	echo "Usage:"
	echo "$0 [ -s ] [ -c <cmdline_file> ] -h <target_host>"
	echo "    -s update stable (KERN_A) partition"
	echo "    -c <cmdline_file> file containing kernel command line"
	echo "       [default: loglevel=7 root=/dev/nvme0n1p3 rootwait rw]"
	echo "    <target_host> target machine to update"
	echo ""
	echo "    Updates <target_host> with the kernel and modules from"
	echo "    the current directory.  Defaults to installing the kernel"
	echo "    into the KERN_B partition and marking it as 'boot once'"
	exit 1
}

function error
{
	echo "ERROR: $1"
	exit 1
}

function default_cmdline
{
	def_cmdline=`mktemp /tmp/push_kernel_cmdline-XXXX`
	cat <<EOF > $def_cmdline
loglevel=7
root=/dev/nvme0n1p3
rootwait
rw
EOF
}

function copy_modules
{
	dir=`mktemp -d /tmp/push_kernel-XXXX`
	echo -n "installing modules..."
	make INSTALL_MOD_PATH=$dir INSTALL_MOD_STRIP=1 modules_install >& /dev/null
	echo "done."
	echo -n "copying modules..."
	(cd $dir/lib/modules; rsync -a * root@$target_host:/lib/modules/)
	echo "done."
	if [ $? -ne 0 ] ; then
		return -1
	fi
	rm -rf $dir
	return 0
}

function build_image
{
	img_file=`mktemp /tmp/push_kernel_img-XXXX`
	vbutil_kernel --pack $img_file \
		      --keyblock /usr/share/vboot/devkeys/kernel.keyblock \
		      --signprivate /usr/share/vboot/devkeys/kernel_data_key.vbprivk \
		      --version 1 \
		      --config "$cmdline_file" \
		      --bootloader ~/bootstub.efi \
		      --vmlinuz arch/x86/boot/bzImage \
		      --arch x86_64
	if [ ! -f $img_file ]; then
		return -1
	fi
	if [ "$def_cmdline" != "" ]; then
		rm $def_cmdline
	fi
	return $?
}
	
function copy_image
{
	# Setup new kernel to boot from partition B
	echo -n "copying kernel image..."
	scp $img_file root@$target_host:/tmp >& /dev/null
	echo "done."
	if [ $? -ne 0 ] ; then
		echo "failed to copy kernel image"
		exit 0
	fi
	rm $img_file
	return 0
}

function install_image
{
	echo -n "installing kernel to partition $partition..."
	ssh root@$target_host dd if=$img_file of=/dev/nvme0n1p$partition bs=4K >& /dev/null
	echo "done."
	if [ $? -ne 0 ] ; then
		echo "failed to write kernel image on target"
		exit 0
	fi
}

function set_boot_once
{
	if [ $partition -eq 2 ] ; then
		echo -n "setting boot flags on KERN_A..."
		ssh root@$target_host cgpt add -i 2 -S 1 -T 1 -P14 /dev/nvme0n1
	else
		echo -n "setting boot once flag for KERN_B..."
		ssh root@$target_host cgpt add -i 4 -S 0 -T 1 -P15 /dev/nvme0n1
	fi
	if [ $? -ne 0 ] ; then
		echo "failed to set new kernel image active on target"
		exit 0
	fi
	echo "done."
}

optstr="sc:h:"

while getopts $optstr opt; do
	case $opt in
		's')
			partition=2
			echo "updating stable kernel (KERN_A)"
			;;
		'c')
			cmdline_file=$OPTARG
			echo "using $cmdline_file for boot command line"
			;;
		'h')
			target_host=$OPTARG
			echo "sending kernel to $target_host"
			;;
		*)
			usage
			;;
	esac
done

if [ "$target_host" == "" ]; then
	usage
fi

if [ "$cmdline_file" == "" ]; then
	default_cmdline
	cmdline_file=$def_cmdline
	echo "using default command line"
fi

if [ ! -f arch/x86/boot/bzImage ] ; then
	echo "run from a linux tree with the kernel built"
	exit 1
fi

if [ ! -x "$(command -v vbutil_kernel)" ]; then
	echo "vbutil_kernel not installed, install vboot-kernel-utils"
	exit 1
fi

copy_modules
if [ $? -ne 0 ]; then
	echo "module copy failed"
	exit 1
fi

build_image
if [ $? -ne 0 ]; then
	echo "image build failed"
	exit 1
fi

copy_image
if [ $? -ne 0 ]; then
	echo "image copy failed"
	exit 1
fi

install_image
if [ $? -ne 0 ]; then
	echo "image install failed"
	exit 1
fi

set_boot_once
if [ $? -ne 0 ]; then
	echo "failed to set kernel to boot"
	exit 1
fi

echo "success, $target_host is ready to reboot."
exit 0