#!/bin/bash
# main autocluster script
#
# Copyright (C) Andrew Tridgell  2008
# Copyright (C) Martin Schwenke  2008
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 3 of the License, or
# (at your option) any later version.
#   
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#   
# You should have received a copy of the GNU General Public License
# along with this program; if not, see <http://www.gnu.org/licenses/>.

##BEGIN-INSTALLDIR-MAGIC##
# There are better ways of doing this but not if you still want to be
# able to run straight out of a git tree.  :-)
if [ -f "$0" ]; then
    installdir="`dirname \"$0\"`"
else
    autocluster=`which $0`
    installdir="`dirname \"$autocluster\"`"
fi
##END-INSTALLDIR-MAGIC##

####################
# show program usage
usage ()
{
    cat <<EOF
Usage: autocluster [OPTION] ... <COMMAND>
  options:
     -c <file>                   specify config file (default is "config")
EOF

    local releases=$(echo $(cd $installdir/releases && ls))
    releases="\"${releases// /\", \"}\""

    usage_smart_display \
	defconf "WITH_RELEASE" "" \
	"<string>" "specify kickstart and yum templates using a release version string.  Possible values are: ${releases}."

cat <<EOF
     -x                          enable script debugging
     --dump                      dump config settings and exit

  configuration options:
EOF

    usage_config_options

    cat <<EOF

  commands:
     create base
           create a base image

     create cluster CLUSTERNAME
           create a full cluster

     create node CLUSTERNAME NODENUMBER
           create one cluster node

     create tsm CLUSTERNAME
           create a TSM server node

     mount DISK
           mount a qemu disk on mnt/

     unmount
           unmount a qemu disk from mnt/

     bootbase
           boot the base image

     testproxy
           test your proxy setup
EOF
    exit 1
}

###############################
# common node creation stuff
create_node_common() {
    local node="$1"

    local nodenumber
    local template_file
    if [ "$node" = "tsm" ] ; then
	nodenumber=0
	NAME="${CLUSTER}tsm"
	template_file=$TSM_TEMPLATE
	echo "Creating TSM cluster node $NAME"
    else
	nodenumber=$node
	NAME="${CLUSTER}n${nodenumber}"
	template_file=$NODE_TEMPLATE
	echo "Creating cluster node $NAME"
    fi

    IPNUM=$(($FIRSTIP + $nodenumber))
    DISK="${VIRTBASE}/${CLUSTER}/${NAME}.qcow2"

    mkdir -p $VIRTBASE/$CLUSTER tmp

    echo "Creating the disk"
    rm -f "$DISK"
    qemu-img create -b "$VIRTBASE/$BASENAME.img" -f qcow2 "$DISK"

    mount_disk $DISK
    setup_base
    setup_network
    unmount_disk

    set_macaddrs $CLUSTER $nodenumber
    UUID=`uuidgen`
    
    echo "Creating $NAME.xml"
    substitute_vars $template_file tmp/$NAME.xml
    
    # install the XML file
    $VIRSH undefine $NAME > /dev/null 2>&1 || true
    $VIRSH define tmp/$NAME.xml
}

###############################
# create a single node
create_node() {
    CLUSTER="$1"
    local nodenumber="$2"

    # first node might need more memory for SoFS GUI
    if [ "$WITH_SOFS_GUI" = 1 -a $nodenumber = 1 -a $GUIMEM -gt $MEM ]; then
	NODEMEM=$GUIMEM
    else
	NODEMEM=$MEM
    fi

    create_node_common "$nodenumber"
}

###############################
# create a TSM node
create_tsm() {
    CLUSTER="$1"

    # This must be before create_node_common.
    TSMDISK="${VIRTBASE}/${CLUSTER}/tsmstorage.qcow2"

    create_node_common "tsm"

    echo "Creating tsm disk"
    qemu-img create -f qcow2 "$TSMDISK" $TSMDISKSIZE
}

###############################
# create a whole cluster
create_cluster() {
    CLUSTER="$1"

    mkdir -p $VIRTBASE/$CLUSTER tmp
    mkdir -p $KVMLOG

    echo "Creating 3 shared disks"
    for i in `seq 1 3`; do
	# setup a nice ID at the start of the disk
	qemu-img create -f raw $VIRTBASE/$CLUSTER/shared$i $SHAREDDISKSIZE
	echo "SOFS-`uuidgen`" > tmp/diskid
	dd if=tmp/diskid of=$VIRTBASE/$CLUSTER/shared$i conv=notrunc bs=1 > /dev/null 2>&1
    done

    echo "Creating $NUMNODES base nodes"
    for i in `seq 1 $NUMNODES`; do
	create_node "$CLUSTER" $i
    done

    if [ $WITH_TSM_NODE -eq 1 ]; then
      echo "Creating TSM server node"
      create_tsm "$CLUSTER"
    fi

    echo "# autocluster $CLUSTER" > hosts.$CLUSTER
    [ $WITH_TSM_NODE -eq 1 ] && {
	echo "$IPBASE.0.$FIRSTIP ${CLUSTER}tsm.$LOWDOMAIN ${CLUSTER}tsm" >> hosts.$CLUSTER
    }
    for i in `seq 1 $NUMNODES`; do
	echo "$IPBASE.0.`expr $FIRSTIP + $i` ${CLUSTER}n$i.$LOWDOMAIN ${CLUSTER}n$i" >> hosts.$CLUSTER
    done
    echo >> hosts.$CLUSTER

    echo "Cluster $CLUSTER created"
    echo "You may want to add this to your /etc/hosts file:"
    cat hosts.$CLUSTER

    echo
}

###############################
# test the proxy setup
test_proxy() {
    export http_proxy=$WEBPROXY
    wget -O /dev/null $INSTALL_SERVER || {
	echo "Your WEBPROXY setting \"$WEBPROXY\" is not working"
	exit 1
    }
    echo "Proxy OK"
}

###################
# create base image
create_base() {

    NAME="$BASENAME"
    DISK="$VIRTBASE/$NAME.img"

    mkdir -p $KVMLOG

    echo "Testing WEBPROXY $WEBPROXY"
    test_proxy

    echo "Creating the disk"
    qemu-img create -f $BASE_FORMAT "$DISK" $DISKSIZE

    rm -rf tmp
    mkdir -p mnt tmp tmp/ISO

    setup_timezone
    setup_indirects

    echo "Creating kickstart file from template"
    substitute_vars "$KICKSTART" "tmp/ks.cfg"

    if [ $INSTALLKEY = "--skip" ]; then
	cat <<EOF
--------------------------------------------------------------------------------------
WARNING: You have not entered an install key. Some RHEL packages will not be installed.

Please enter a valid RHEL install key in your config file like this:

  INSTALLKEY="1234-5678-0123-4567"

The install will continue without an install key in 5 seconds
--------------------------------------------------------------------------------------
EOF
	sleep 5
    fi
    
    echo "Creating kickstart floppy"
    dd if=/dev/zero of=tmp/floppy.img bs=1024 count=1440
    mkdosfs tmp/floppy.img
    mount -o loop -t msdos tmp/floppy.img mnt
    cp tmp/ks.cfg mnt
    mount -o loop,ro $ISO tmp/ISO
    
    echo "Setting up bootloader"
    cp tmp/ISO/isolinux/isolinux.bin tmp
    cp tmp/ISO/isolinux/vmlinuz tmp
    cp tmp/ISO/isolinux/initrd.img tmp
    umount tmp/ISO
    umount mnt

    UUID=`uuidgen`

    substitute_vars $INSTALL_TEMPLATE tmp/$NAME.xml

    rm -f $KVMLOG/serial.$NAME

    # boot the install CD
    $VIRSH create tmp/$NAME.xml

    echo "Waiting for install to start"
    sleep 2
    
    # wait for the install to finish
    if ! waitfor $KVMLOG/serial.$NAME "you may safely reboot your system" 3600; then
	$VIRSH destroy $NAME
	echo "Failed to create base image $DISK"
	exit 1
    fi
    
    $VIRSH destroy $NAME

    ls -l $DISK
    cat <<EOF

Install finished, base image $DISK created

You may wish to run
   chattr +i $DISK
To ensure that this image does not change

Note that the root password has been set to $ROOTPASSWORD

EOF
}

###############################
# boot the base disk
boot_base() {
    CLUSTER="$1"

    NAME="$BASENAME"
    DISK="$VIRTBASE/$NAME.img"

    rm -rf tmp
    mkdir -p tmp

    IPNUM=$FIRSTIP
    CLUSTER="base"

    mount_disk $DISK
    setup_base
    unmount_disk

    UUID=`uuidgen`
    
    echo "Creating $NAME.xml"
    substitute_vars $BOOT_TEMPLATE tmp/$NAME.xml
    
    # boot the base system
    $VIRSH create tmp/$NAME.xml
}

######################################################################

# various functions...

# Set some MAC address variables based on a hash of the cluster name
# plus the node number and each adapter number.
set_macaddrs () {
    local cname="$1"
    local nodenumber="$2"

    local md5=$(echo $cname | md5sum)
    local nh=$(printf "%02x" $nodenumber)
    local mac_prefix="02:${md5:0:2}:${md5:2:2}:00:${nh}:"

    MAC1="${mac_prefix}01"
    MAC2="${mac_prefix}02"
    MAC3="${mac_prefix}03"
    MAC4="${mac_prefix}04"
    MAC5="${mac_prefix}05"
    MAC6="${mac_prefix}06"
}

# mount a qemu image via nbd
connect_nbd() {    
    echo "Connecting nbd to $1"
    mkdir -p mnt
    modprobe nbd
    killall -9 -q $QEMU_NBD || true
    $QEMU_NBD -p 1300 $1 &
    sleep 1
    [ -r $NBD_DEVICE ] || {
	mknod $NBD_DEVICE b 43 0
    }
    umount mnt 2> /dev/null || true
    nbd-client -d $NBD_DEVICE > /dev/null 2>&1 || true
    killall -9 -q nbd-client || true
    nbd-client localhost 1300 $NBD_DEVICE > /dev/null 2>&1 || true &
    sleep 1
}

# disconnect nbd
disconnect_nbd() {
    echo "Disconnecting nbd"
    sync; sync
    nbd-client -d $NBD_DEVICE > /dev/null 2>&1 || true
    killall -9 -q nbd-client || true
    killall -q $QEMU_NBD || true
}

# mount a qemu image via nbd
mount_disk() {    
    connect_nbd $1
    echo "Mounting disk $1"
    mount_ok=0
    for i in `seq 1 5`; do
	mount -o offset=32256 $NBD_DEVICE mnt && {
	    mount_ok=1
	    break
	}
	umount mnt 2>/dev/null || true
	sleep 1
    done
    [ $mount_ok = 1 ] || {
	echo "Failed to mount $1"
	exit 1
    }
    [ -d mnt/root ] || {
	echo "Mounted directory does not look like a root filesystem"
	ls -latr mnt
	exit 1
    }
}

# unmount a qemu image
unmount_disk() {
    echo "Unmounting disk"
    sync; sync;
    umount mnt || umount mnt || true
    disconnect_nbd
}

# setup the files from $BASE_TEMPLATES/, substituting any variables
# based on the config
setup_base() {
    umask 022
    setup_indirects
    echo "Copy base files"
    for f in `cd $BASE_TEMPLATES && find . \! -name '*~'`; do
	if [ -d "$BASE_TEMPLATES/$f" ]; then
	    mkdir -p mnt/"$f"
	else 
	    substitute_vars "$BASE_TEMPLATES/$f" "mnt/$f"
	fi
	chmod --reference="$BASE_TEMPLATES/$f" "mnt/$f"
    done
    # this is needed as git doesn't store file permissions other
    # than execute
    chmod 600 mnt/etc/ssh/*key mnt/root/.ssh/*
    chmod 700 mnt/etc/ssh mnt/root/.ssh mnt/root
    if [ -r "$HOME/.ssh/id_rsa.pub" ]; then
       echo "Adding $HOME/.ssh/id_rsa.pub to ssh authorized_keys"
       cat "$HOME/.ssh/id_rsa.pub" >> mnt/root/.ssh/authorized_keys
    fi
    if [ -r "$HOME/.ssh/id_dsa.pub" ]; then
       echo "Adding $HOME/.ssh/id_dsa.pub to ssh authorized_keys"
       cat "$HOME/.ssh/id_dsa.pub" >> mnt/root/.ssh/authorized_keys
    fi
    echo "Adjusting grub.conf"
    local o="$EXTRA_KERNEL_OPTIONS"
    sed -e "s/console=ttyS0,19200/console=ttyS0,115200/"  \
	-e "s/ nodmraid//" -e "s/ nompath//"  \
	-e "s/quiet/divider=10${o:+ }${o}/g" mnt/boot/grub/grub.conf -i.org
}

setup_indirects() {
    # Order is very important here, since postinstall.sh can use
    # @@YUM_REPOS@@.

    substitute_vars $YUM_TEMPLATE tmp/SOFS.repo
    YUM_REPOS="`cat tmp/SOFS.repo`"

    substitute_vars $POSTINSTALL_TEMPLATE tmp/postinstall.sh
    EXTRA_POSTINSTALL="`cat tmp/postinstall.sh`"

    substitute_vars "$installdir/templates/force-net.sh" tmp/force-net.sh
    AUTOCLUSTER_POSTINSTALL="`cat tmp/force-net.sh`"
}

# setup various networking components
setup_network() {
    echo "Setting up networks"

    for i in `seq 1 $NUMNODES`; do
	echo "$IPBASE.0.`expr $FIRSTIP + $i` ${CLUSTER}n$i.$LOWDOMAIN ${CLUSTER}n$i" >> mnt/etc/hosts
    done

    echo "Setting up /etc/ctdb/nodes"
    mkdir -p mnt/etc/ctdb
    rm -f mnt/etc/ctdb/nodes
    for i in `seq 1 $NUMNODES`; do
	echo "$IPBASE.0.`expr $FIRSTIP + $i`" >> mnt/etc/ctdb/nodes
    done

    [ "$WEBPROXY" = "" ] || {
	echo "export http_proxy=$WEBPROXY" >> mnt/etc/bashrc
    }

    echo "Enabling nfs mount of $NFSSHARE"
    mkdir -p mnt/root/SOFS
    echo "$NFSSHARE /root/SOFS nfs intr" >> mnt/etc/fstab

    mkdir -p mnt/etc/yum.repos.d
    substitute_vars "$YUM_TEMPLATE" mnt/etc/yum.repos.d/SOFS.repo
}

setup_timezone() {
    [ -z "$TIMEZONE" ] && {
	[ -r /etc/timezone ] && {
	    TIMEZONE=`cat /etc/timezone`
	}
	[ -r /etc/sysconfig/clock ] && {
	    . /etc/sysconfig/clock
	    TIMEZONE="$ZONE"
	}
    }
    if [ "$TIMEZONE" = "" ]; then
	echo "Unable to determine TIMEZONE - please set in config"
	exit 1
    fi
}

# substite a set of variables of the form @@XX@@ for the shell
# variables $XX in a file
substitute_vars() {(
	infile="$1"
	outfile="$2"
	delimiter="@"
	delim2="$delimiter$delimiter"

	# get the list of variables used in the template
	VARS=`tr -cs "A-Z0-9_$delimiter" '\012' < "$infile" | 
              sort -u |
	      egrep "^$delim2.*$delim2$" |
	      cut -d$delimiter -f3`

	rm -f sed.$$
	touch sed.$$
	for v in $VARS; do
	    # variable variables are fun .....
	    [ "${!v+x}" ] || {
		echo "ERROR: No substitution given for ${delim2}$v${delim2} in $infile"
		rm -f sed.$$
		exit 1
	    }
	    s=${!v}
	    # escape some pesky chars
	    s=${s//
/\\n}
	    s=${s//#/\\#}
	    s=${s//&/\\&}
	    echo "s#@@$v@@#$s#g" >> sed.$$
	done

	sed -f sed.$$ < "$infile" > "$outfile" || exit 1
	rm -f sed.$$
)}


check_command() {
    if which $1 > /dev/null; then
	# echo "$1 is installed: OK"
	:
    else
	echo "Please install $1 to continue"
	exit 1
    fi
}

# Set a variable if it isn't already set.  This allows environment
# variables to override default config settings.
defconf() {
    local v="$1"
    local e="$2"

    [ "${!v+x}" ] || eval "$v=\"$e\""
}

# Print the list of config variables defined in config.default.
get_config_options() {( # sub-shell for local declaration of defconf()
	local options=
	defconf() { options="$options $1" ; }
	. "$installdir/config.default"
	echo $options
)}

# Produce a list of long options, suitable for use with getopt, that
# represent the config variables defined in config.default.
getopt_config_options() {
    local x=$(get_config_options | tr 'A-Z_' 'a-z-')
    echo "${x// /:,}:"
}

# Unconditionally set the config variable associated with the given
# long option.
setconf_longopt() {
    local longopt="$1"
    local e="$2"

    local v=$(echo "${longopt#--}" | tr 'a-z-' 'A-Z_')
    # unset so defconf will set it
    eval "unset $v"
    defconf "$v" "$e"
}

# Dump all of the current config variables.
dump_config() {
    local o
    for o in $(get_config_options) ; do
	echo "${o}=\"${!o}\""
    done
    exit 0
}

# Print text assuming it starts after other text in $startcol and
# needs to wrap before $fillcol.  Subsequent lines start at $startcol.
# Long "words" will extend past $fillcol.
fill_text() {
    local startcol="$1"
    local fillcol="$2"
    local text="$3"

    local width=$(($fillcol - $startcol))
    [ $width -lt 0 ] && width=$((78 - $startcol))

    local out=""

    local padding=$(printf "\n%${startcol}s" " ")

    while [ -n "$text" ] ; do
	local orig="$text"

	# If we already have output then arrange padding on the next line.
	[ -n "$out" ] && out="${out}${padding}"

	# Break the text at $width.
	out="${out}${text:0:${width}}"
	text="${text:${width}}"

	# If we have left over text then the line break may be ugly,
	# so let's check and try to break it on a space.
	if [ -n "$text" ] ; then
	    if [ "${text:0:1}" != " " -a "${text: -1:1}" != " " ] ; then
		# We didn't break on a space.  Arrange for the
		# beginning of the broken "word" to appear on the next
		# line but not if it will make us loop infinitely.
		if [ "${orig}" != "${out##* }${text}" ] ; then
		    text="${out##* }${text}"
		    out="${out% *}"
		else
		    # Hmmm, doing that would make us loop, so add the
		    # rest of the word from the remainder of the text
		    # to this line and let it extend past $fillcol.
		    out="${out}${text%% *}"
		    if [ "${text# *}" != "$text" ] ; then
			# Remember the text after the next space for next time.
			text="${text# *}"
		    else
			# No text after next space.
			text=""
		    fi
		fi
	    else
		# We broke on a space.  If it will be at the beginning
		# of the next line then remove it.
		text="${text# }"
	    fi
	fi
    done

    echo "$out"
}

# Display usage text, trying these approaches in order.
# 1. See if it all fits on one line before $fillcol.
# 2. See if splitting before the default value and indenting it
#    to $startcol means that nothing passes $fillcol.
# 3. Treat the message and default value as a string and just us fill_text()
#    to format it. 
usage_display_text() {
    local startcol="$1"
    local fillcol="$2"
    local desc="$3"
    local default="$4"
    
    local width=$(($fillcol - $startcol))

    # To avoid clutter, only quote values that need it.
    #local q=
    #[ -z "$default" -o "$(echo $default)" = "${default// /}" ] || q="\""
    default="(default \"$default\")"

    if [ $((${#desc} + 1 + ${#default})) -le $width ] ; then
	echo "${desc} ${default}"
    else
	local padding=$(printf "%${startcol}s" " ")

	if [ ${#desc} -lt $width -a ${#default} -lt $width ] ; then
	    echo "$desc"
	    echo "${padding}${default}"
	else
	    fill_text $startcol $fillcol "${desc} ${default}"
	fi
    fi
}

# Display usage information for long config options.
usage_smart_display () {( # sub-shell for local declaration of defconf()
	local def_fillcol=78
	local fillcol=$def_fillcol
	local rows=$(stty size | sed -e 's@.* @@')
	[ -n "$rows" ] && fillcol=$(($rows - 2))

	local startcol=33

	# We need to have at least one column... negative is also a problem.
	[ $fillcol -le $startcol ] && fillcol=$def_fillcol

	defconf() {
	    local local longopt=$(echo "$1" | tr 'A-Z_' 'a-z-')

	    printf "     --%-25s " "${longopt}=${3}"

	    usage_display_text $startcol $fillcol "$4" "$2"
	}

	"$@"
)}


# Display usage information for long config options.
usage_config_options(){
    usage_smart_display . "$installdir/config.default"
}

######################################################################

. "$installdir/config.default"

############################
# parse command line options
long_opts=$(getopt_config_options)
getopt_output=$(getopt -n autocluster -o "c:xh" -l help,dump,with-release -l "$long_opts" -- "$@")
[ $? != 0 ] && usage

use_default_config=true

# We do 2 passes of the options.  The first time we just handle usage
# and check whether -c is being used.
eval set -- "$getopt_output"
while true ; do
    case "$1" in
	-c) shift 2 ; use_default_config=false ;;
	--) shift ; break ;;
	--with-release) shift 2 ;; # Don't set use_default_config=false!!!
	--dump|-x) shift ;;
	-h|--help) usage ;; # Usage should be shown here for real defaults.
	--*) shift 2 ;; # Assume other long opts are valid and take an arg.
	*) usage ;; # shouldn't happen, so this is reasonable.
    esac
done

config="./config"
$use_default_config && [ -r "$config" ] && . "$config"

eval set -- "$getopt_output"

while true ; do
    case "$1" in
	-c) . "`dirname $2`/$2" ; shift 2 ;;
	--with-release)
	    # This simply loads an extra config file from $installdir/releases
	    f="${installdir}/releases/$2"
	    if [ -r "$f" ] ; then
		. "$f"
	    else
		echo "Unknown release \"$2\" specified to --with-release"
		exit 1
	    fi
	    shift 2
	    ;;
	-x) set -x; shift ;;
	--dump) dump_config ;;
	--) shift ; break ;;
	-h|--help) usage ;; # Redundant.
	--*)
	    # Putting --opt1|opt2|... into a variable and having case
	    # match against it as a pattern doesn't work.  The | is
	    # part of shell syntax, so we need to do this.  Look away
	    # now to stop your eyes from bleeding! :-)
	    x=",${long_opts}" # Now each option is surrounded by , and :
	    if [ "$x" != "${x#*,${1#--}:}" ] ; then
		# Our option, $1, surrounded by , and : was in $x, so is legal.
		setconf_longopt "$1" "$2"; shift 2
	    else
		usage
	    fi
	    ;;
	*) usage ;; # shouldn't happen, so this is reasonable.
    esac
done

# catch errors
set -e
set -E
trap 'es=$?; 
      echo ERROR: failed in function \"${FUNCNAME}\" at line ${LINENO} of ${BASH_SOURCE[0]} with code $es; 
      exit $es' ERR

LOWDOMAIN=`echo $DOMAIN | tr A-Z a-z`

# check for needed programs 
check_command nbd-client
check_command expect
check_command $QEMU_NBD

[ $# -lt 1 ] && usage

command="$1"
shift

case $command in
    create)
	type=$1
	shift
	case $type in
	    base)
		[ $# != 0 ] && usage
		create_base;
		;;
	    cluster)
		[ $# != 1 ] && usage
		create_cluster "$1";
		;;
	    node)
		[ $# != 2 ] && usage
		create_node "$1" "$2";
		;;
	    tsm)
		[ $# != 1 ] && usage
		create_tsm "$1";
		;;
	    *)
		usage;
		;;
	esac
	;;
    mount)
	[ $# != 1 ] && usage
	mount_disk "$1"
	;;
    unmount)
	[ $# != 0 ] && usage
	unmount_disk
	;;
    bootbase)
	boot_base;
	;;
    testproxy)
	test_proxy;
	;;
    *)
	usage;
	;;
esac