#!/bin/bash
#
# Steeltoe command processor

fail()
{
	# Print error message on $1 to stderr
	echo "$1" >&2
	# then exit with error code $2 or 1
	exit ${2:-1}
}

# require: make sure specified config options are set
require()
{
	for vn in $*; do case $vn in
		RSH) RSH=`gxpp "/steeltoe/config/remoteshell" $STXML` || fail "Config option remoteshell not found";;
		RCP) RCP=`gxpp "/steeltoe/config/remotecopy" $STXML` || fail "Config option remotecopy not found";;
		TFTPROOT) TFTPROOT=`gxpp "/steeltoe/config/tftproot" $STXML` || fail "Config option tftproot not found"
			  [ -d "$TFTPROOT" -a -w "$TFTPROOT" ] || fail "Directory $TFTPROOT does not exist or is not writeable";;
		TREESERVER) TREESERVER=`gxpp "/steeltoe/config/treeserver" $STXML` || fail "Config option treeserver not found";;
		KSBASE) KSBASE=`gxpp "/steeltoe/config/kickstartbase" $STXML` || fail "Config option kickstartbase not found";;
		YABOOTETC) YABOOTETC=`gxpp "/steeltoe/config/yabootetc" $STXML` || fail "Config option yabootetc not found"
			   [ -d "$YABOOTETC" -a -w "$YABOOTETC" ] || fail "Directory $YABOOTETC does not exist or is not writeable";;
		*) fail "Unknown config requirement '$vn'"
	esac; done
}

expand_groups() {
	for target in $*; do
		grouplist=$(gxpp "//group[@name='$target']" $STXML)
		if [ -z "$grouplist" ]; then
			echo $target
		else 
			for g in $grouplist; do echo $g; done
		fi
	done
}

# Set global config variables
if [ -z "$STEELTOE_ETC" ]; then
	STEELTOE_ETC="/etc/steeltoe"
fi
STXML="$STEELTOE_ETC/steeltoe.xml";
if [ ! -f $STXML ]; then
	fail "Steel Toe config file '$STXML' not found"
fi

XSLTOPTS="--xinclude"
# Make the default install method http
instmethod="http"

check_instmethod()
{
	instmethods="http ftp nfs"
	if ! echo "$instmethods" | grep -q $1; then
		echo "Install method '$1' unknown.  Supported methods: $instmethods" >&2
		exit 1
	fi
}

# Make the default kickstart template /etc/steeltoe/kickstart.xsl
kstemplate="$STEELTOE_ETC/kickstart.xsl"

check_kstemplate()
{
  if [ ! -f $1 ]; then
    echo "Kickstart template '$1' does not exist" >&2
    exit 1
  fi
}

global_usage()
{
	cat >&2 << __EOT__ 
Usage: $0 <cmd> <options>
__EOT__
	for i in $COMMANDS; do
		${i}_usage >&2
	done
}

COMMANDS="$COMMANDS gendhcpd"
gendhcpd_usage()
{
	cat << __EOT__
	gendhcpd [-o outputfile ]
		Generate ISC DHCPD configuration snippet
__EOT__
}

gendhcpd_do()
{
	if [ -n "$outputfile" ]; then
		XSLTOPTS="$XSLTOPTS -o $outputfile"
	fi
	xsltproc $XSLTOPTS $STEELTOE_ETC/dhcpd.xsl $STXML
}

COMMANDS="$COMMANDS genks"
genks_usage()
{
	cat << __EOT__
	genks [-k kickstart.xsl][-m http|ftp|nfs][-o outputfile ] host tree
		Generate a Kickstart config file for host to install tree
__EOT__
}

genks_do()
{
	if [ -z "$1" -o -z "$2" ]; then
		genks_usage >&2
		exit 1
	fi
	if [ -n "$outputfile" ]; then
		XSLTOPTS="$XSLTOPTS -o $outputfile"
	fi
	if [ -n "$instmethod" ]; then
		METHOD="--stringparam method $instmethod"
	fi
	xsltproc $XSLTOPTS $METHOD --stringparam hostname "$1" --stringparam treename "$2" $kstemplate $STXML
}

COMMANDS="$COMMANDS genpxe"
genpxe_usage()
{
	cat << __EOT__
	genpxe [-o outputfile ] host tree
		Generate a PXE config file for host to install tree
__EOT__
}

genpxe_do()
{
	if [ -z "$1" -o -z "$2" ]; then
		genpxe_usage >&2
		exit 1
	fi
	if [ -n "$outputfile" ]; then
		XSLTOPTS="$XSLTOPTS -o $outputfile"
	fi
	xsltproc $XSLTOPTS --stringparam hostname "$1" --stringparam treename "$2" $STEELTOE_ETC/pxeconf.xsl $STXML
}

COMMANDS="$COMMANDS genelilo"
genelilo_usage()
{
	cat << __EOT__
	genelilo [-o outputfile ] host tree
		Generate a ELILO config file for host to install tree
__EOT__
}

genelilo_do()
{
	if [ -z "$1" -o -z "$2" ]; then
		genelilo_usage >&2
		exit 1
	fi
	if [ -n "$outputfile" ]; then
		XSLTOPTS="$XSLTOPTS -o $outputfile"
	fi
	xsltproc $XSLTOPTS --stringparam hostname "$1" --stringparam treename "$2" $STEELTOE_ETC/eliloconf.xsl $STXML
}

COMMANDS="$COMMANDS genyaboot"
genyaboot_usage()
{
	cat << __EOT__
	genyaboot [-o outputfile ] host tree
		Generate a yaboot config file for host to install tree
__EOT__
}

genyaboot_do()
{
	if [ -z "$1" -o -z "$2" ]; then
		genyaboot_usage >&2
		exit 1
	fi
	if [ -n "$outputfile" ]; then
		XSLTOPTS="$XSLTOPTS -o $outputfile"
	fi
	xsltproc $XSLTOPTS --stringparam hostname "$1" --stringparam treename "$2" $STEELTOE_ETC/yaboot.xsl $STXML
}

COMMANDS="install $COMMANDS"
install_usage()
{
	cat << __EOT__
	install [options] tree host|group [host|group ...]
		Install tree on hosts
		-a ARCH	override the architecture of a host
		-k FILE	specify an alternate kickstart template to use.
		-m meth	install method (http, ftp, or nfs; default http)
		-r	reboot now
__EOT__
}

install_do()
{
	if [ -z "$1" -o -z "$2" ]; then
		install_usage >&2
		exit 1
	fi
	require TFTPROOT RSH
	treename="$1"
	shift
	hostlist=`expand_groups $*`

	for host in $hostlist; do
		echo "Configuring $host to install $treename"
		tftp_do $host $treename
		if ping -q -c 1 -w 3 $host > /dev/null && $RSH -l root $host /bin/true; then
			autoks_do $host $treename
		else
			echo "$host needs to be manually PXE booted to install $treename"
		fi
	done
}

COMMANDS="$COMMANDS tftp"
tftp_usage()
{
	cat << __EOT__
	tftp [-m http|ftp|nfs][-a ARCH] host tree
		Prepare host's TFTP directory for PXE installing tree.
		-a ARCH	override the architecture of a host
		-m	install method (default http)
__EOT__
}

tftp_do()
{
	if [ -z "$1" -o -z "$2" ]; then
		tftp_usage >&2
		exit 1
	fi
	require TFTPROOT RSH
	host="$1"
	treename="$2"
	hostdir="$TFTPROOT/$host"

	# Try to lookup something under the host so we catch a non-existant host early.
	hostarch=`gxpp "//host[@name='$host']/arch" $STXML` || fail "Could not look up architecture of host '$host'"

	# Check to see if we're overriding the arch on the command line
	if [ -n "$forcearch" ]; then
		hostarch=$forcearch
	fi

	# First lookup the tree information to make sure it exists.
	treebase=`gxpp "//tree[@name='$treename' and meta/arch='$hostarch']/path" $STXML` || fail "Could not find tree '$treename' for host '$host'"
	# Next figure out which bootloader we need to configure
	loadername=`gxpp "//host[@name='$host']/bootloader" $STXML` || fail "Boot loader not specified for host '$host'"
	# Now find the source for the bootloader link
	loadersrc=`gxpp "/steeltoe/config/bootloader/$loadername" $STXML` || fail "Boot loader source is not set for '$loadername'"
	# Make sure the host's TFTP directory exists
	[ -d "$hostdir" ] || mkdir -p $hostdir
	# Link in the bootloader specified in the config file
	[ -f "$hostdir/$loadername" ] || ln -f $loadersrc $hostdir/$loadername
	# Generate the kickstart config
	genks_do $host $treename > $hostdir/ks.cfg
	# Handle bootloader specific configuration needs and link in the kernel and initrd
	case $loadername in
	  pxelinux.0) genpxe_do $host $treename > $hostdir/default
			[ -e "$hostdir/pxelinux.cfg" ] || ln -s . $hostdir/pxelinux.cfg
			ln -sf $treebase/images/pxeboot/vmlinuz $hostdir/vmlinuz
			ln -sf $treebase/images/pxeboot/initrd.img $hostdir/initrd.img
			;;
	   elilo.efi) genelilo_do $host $treename > $hostdir/elilo.conf
			ln -sf $treebase/images/pxeboot/vmlinuz $hostdir/vmlinuz
			ln -sf $treebase/images/pxeboot/initrd.img $hostdir/initrd.img
			;;
	      yaboot) genyaboot_do $host $treename > $hostdir/yaboot.conf
			require YABOOTETC
			etcfn="01-`gxpp "//host[@name='$host']/macaddr" $STXML | tr [A-Z]: [a-z]-`"
			ln -f $hostdir/yaboot.conf $YABOOTETC/$etcfn
			# kernel and ramdisk locations differ between RHEL4 and RHEL5
			for subdir in ppc/chrp ppc/ppc64; do if [ -f $treebase/$subdir/vmlinuz ]; then
				ln -sf $treebase/$subdir/vmlinuz $hostdir/vmlinuz
				ln -sf $treebase/$subdir/ramdisk.image.gz $hostdir/initrd.img
				break
			fi; done;;
	esac
}

COMMANDS="$COMMANDS autoks"
autoks_usage()
{
	cat << __EOT__
	autoks [-k kickstart.xsl][-r] host tree
		Copy installer to host and configure bootloader to boot it.
		-k FILE	Specify alternate kickstart template
		-r	reboot now
__EOT__
}

autoks_do()
{
	require RSH RCP TFTPROOT
	if [ -z "$1" -o -z "$2" ]; then
		autoks_usage >&2
		exit 1
	fi
	host="$1"
	treename="$2"

	$RSH -l root $host /bin/true || fail "Host $host must be up to auto-kickstart"

	# Try to lookup something under the host so we catch a non-existant host early.
	hostarch=`gxpp "//host[@name='$host']/arch" $STXML` || fail "Could not look up architecture of host '$host'"

	# Figure out where to copy the vmlinuz and initrd on the host
	if [ "$hostarch" = "ia64" ]; then
		bootdir="/boot/efi/EFI/redhat"
	else
		bootdir="/boot"
	fi

	# First lookup the tree information to make sure it exists.
	ksargs=`xsltproc $XSLTOPTS --stringparam hostname "$host" --stringparam treename "$treename" $STEELTOE_ETC/autoks-args.xsl $STXML` || fail "Could not generate kickstart arguments for $host and $treename"
	hostdir="$TFTPROOT/$host"
	# Make sure the host's TFTP directory is populated
	[ -d "$hostdir" ] || fail "No TFTP directory for $host, run tftp command first or install"
	[ -e "$hostdir/vmlinuz" ] || fail "No kernel $host, run tftp command first or install"
	[ -e "$hostdir/initrd.img" ] || fail "No ramdisk for $host, run tftp command first or install"

	# Everything is in place, copy it out to host
	echo "Copying files to $host for auto-kickstart of $treename"
	$RCP $hostdir/vmlinuz root@$host:$bootdir/vmlinuz-steeltoe || fail "Failed to copy vmlinuz"
	$RCP $hostdir/initrd.img root@$host:$bootdir/initrd.img-steeltoe || fail "Failed to copy initrd"

	# Check to see if there is a steeltoe kernel and config on host already so we don't duplicate it.
	$RSH -l root $host "/sbin/grubby --set-default $bootdir/vmlinuz-steeltoe"
	defkernel=`$RSH -l root $host "/sbin/grubby --default-kernel --bad-image-okay"`
	if [ "$defkernel" != "$bootdir/vmlinuz-steeltoe" ]; then
		# There wasn't a bootloader entry already, make one
		$RSH -l root $host "/sbin/grubby --make-default --title='steeltoe' --add-kernel=$bootdir/vmlinuz-steeltoe --initrd=$bootdir/initrd.img-steeltoe --args='$ksargs'" || fail "Failed to configure bootloader"
	fi
	$RSH -l root $host sync
	if [ $reboot ]; then
		echo "Rebooting $host"
		$RSH -l root $host "reboot -fi" > /dev/null 2>&1 &
	fi
}


COMMANDS="$COMMANDS virtinstall"
virtinstall_usage()
{
	cat << __EOT__
	virtinstall tree virthost installhost virtroot_device [virtshare_device [virtshare_device] [...]]
		Install tree on virthost from installhost 
		onto virtroot_device, and configure virt guest 
		profile to use virtshare_device(s).
__EOT__
}

virtinstall_do()
{
	if [[ $# < 4 ]]; then
		virtinstall_usage >&2
		exit 2
	fi

	treename="$1"
	virthost="$2"
	installhost="$3"
	virtroot="$4"
	shift 4
	#virtshare is all the rest...

	# Don't even bother if the virt host is up and running
	if ping -c 1 -W 1 $virthost > /dev/null ; then
		fail "$virthost seems to be running, cannot install on running virt guest"
	fi

	require RSH TFTPROOT TREESERVER KSBASE
	# Let install_do do all the hard work and set up the tftp area.
	tftp_do $virthost $treename 

	virtmac=$(gxpp "//host[@name='$virthost']/macaddr" $STXML) || fail "Unable to find $virthost mac address in $STXML"
	DEFAULT_RAM=1024
	virtram=$(gxpp "//host[@name='$virthost']/ram" $STXML)
	virtram=${virtram:-$DEFAULT_RAM}

	ks="$KSBASE/$virthost/ks.cfg"

	hostarch=$(gxpp "//host[@name='$virthost']/arch" $STXML) || fail "Could not look up architecture of virthost '$virthost'"
	treebase=$(gxpp "//tree[@name='$treename' and meta/arch='$hostarch']/path" $STXML) || fail "Could not find tree '$treename' for virthost '$virthost'"
	location="http://$TREESERVER/$treebase"

	extra="'ks=$ks ro selinux=0'"

	# Rebuild the kickstart file, setting the virtnstall param on the xsltproc cmdline.
	xsltproc --xinclude --stringparam virtinstall true --stringparam method http --stringparam hostname $virthost --stringparam treename $treename $kstemplate $STXML > $TFTPROOT/$virthost/ks.cfg

	# Kick off the virtinstall on the installhost
	$RSH -l root $installhost "virt-install -n $virthost -f $virtroot -r $virtram -m $virtmac -p -x $extra --nographics -l $location --noautoconsole"

	# Add to the Xen config file created to add the share devices
	# Blech, Xen specific stuff follows.  Right now it seems we
	# can only have xvda-xvdg as disks, so limit shared to b-g
	XEN_CONFIGROOT=/etc/xen
	SHARED_DEVICES=(xvdb xvdc xvdd xvde xvdf xvdg)
	i=0
	for vshare in $*; do
		vshareentry="disk.append(\"phy:${vshare},${SHARED_DEVICES[$i]},w!\")"
		echo "Adding $vshare to $virthost"
		$RSH -l root $installhost "echo '$vshareentry' >> $XEN_CONFIGROOT/$virthost"
		(( i = i + 1 ))
		if [[ $i > 6 ]]; then
			echo "Only support 5 shared devices, ignoring the rest"
		fi
	done
}


pxe_reboot() 
{
	host="$1"

	set -o pipefail
	rootdev=`gxpp "//host[@name='$host']/kickstart/rootdisk" $STXML | head -1` || fail "Root disk is not specified"
	echo "Making /dev/$rootdev on $host unbootable"
	$RSH -l root $host dd if=/dev/zero of=/dev/$rootdev bs=1 seek=510 count=2 conv=sync > /dev/null 2>&1
	# Sync
	$RSH -l root $host "sync"
	# reboot
	echo "Rebooting $host"
	$RSH -l root $host "reboot -fin" > /dev/null 2>&1 &
}


efi_reboot()
{
	host="$1"

	# Grab the EFI boot menu entry for Steeltoe
	steeltoe=`$RSH -l root $host "efibootmgr" | awk -F\* '($1 ~ /^Boot/) && ($2 ~ /[Ss]teeltoe/) {print $1}'`

	if [ -z "$steeltoe" ]; then
		fail "Unable to find Steeltoe EFI boot entry on $host"
	fi

	# Set the NextBoot option to the found entry
	echo "Updating EFI on $host to boot to Steeltoe"
	$RSH -l root $host "efibootmgr -n ${steeltoe/Boot/}" || fail "Unable to set NextBoot to Steeltoe on $host"
	# reboot
	echo "Rebooting $host"
	$RSH -l root $host "reboot -fin" > /dev/null 2>&1 &
}

COMMANDS="$COMMANDS list"
list_usage()
{
	cat << __EOT__
	list [-v] <host|group|tree> [substring]
		List out hosts or trees
__EOT__
}

list_do()
{
	class="$1"
	search="$2"
	if [ "$class" == "host" -o "$class" == "tree" -o "$class" == "group" ]; then
		XSLTOPTS="$XSLTOPTS --stringparam class $class"
	elif [ "$class" == "trees" ]; then
		XSLTOPTS="$XSLTOPTS --stringparam class tree"
	elif [ "$class" == "hosts" ]; then
		XSLTOPTS="$XSLTOPTS --stringparam class host"
	elif [ "$class" == "groups" ]; then
		XSLTOPTS="$XSLTOPTS --stringparam class group"
	else
		list_usage >&2
		exit 1
	fi
	if [ -n "$verbose" ]; then
		XSLTOPTS="$XSLTOPTS --param verbose 1"
	fi
	xsltproc $XSLTOPTS $STEELTOE_ETC/list.xsl $STXML
}


# Make sure the first option is a command
cmd="$1"
if [ -z "$cmd" ]; then
	global_usage
	exit 1
fi
if ! echo $COMMANDS | grep -q -- "$cmd"; then
	echo "Unknown command '$cmd'" >&2
	global_usage
	exit 1
fi
shift

# Generate the getopts string based on command
opts="h"
case $cmd in
	gendhcpd) opts="o:$opts";;
	genpxe) opts="o:$opts";;
	genelilo) opts="o:$opts";;
	genyaboot) opts="o:$opts";;
	genks) opts="m:o:k:$opts";;
	install) opts="a:rm:k:$opts";;
	autoks) opts="a:rm:k:$opts";;
	virtinstall) opts="k:$opts";;
	list) opts="v$opts";;
esac

# parse command line options for all subcommands
while getopts "$opts" opt; do
	case $opt in
		a) forcearch="$OPTARG";;
		o) outputfile="$OPTARG";;
		r) reboot=1;;
		v) verbose=1;;
		m) instmethod="$OPTARG"; check_instmethod "$OPTARG";;
		k) kstemplate="$OPTARG"; check_kstemplate "$OPTARG";;
		*) global_usage;;
	esac
done
shift $(($OPTIND - 1))
# Run the command
${cmd}_do $*