#!/bin/bash
# Copyright (C) 2016 Ian Kelling

# 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 2
# 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, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.

# Setup dhcp server to point to tftp server,
# and depending on the type, setup the tftp server.

# usage: $0 TYPE
# default distro is the base debian/fedora type. others are fai & arch.
# for no pxe server, use a no-op like : or true.

set -eE -o pipefail
trap 'echo "$0:$LINENO:error: \"$BASH_COMMAND\" returned $?" >&2' ERR


usage() {
    cat <<EOF
Usage: ${0##*/} [OPTIONS] [HOST TYPE]
Configure dnsmasq pxe server options and fai-chboot if appropriate.

Without HOST TYPE, disable pxe server and fai server.

HOST       A hostname known to the dhcp server, or default for all.
TYPE       One of arch, plain, fai.


-r         Don't redeploy fai config. For example, if there is a different host
           that is mid-install.

-a         Don't setup pxe, just Wait for 2 dhcp acks, then disable the pxe
           server after a delay.  First ack is for pxe boot, 2nd ack is
           for os boot. Sometimes on debian, there is a 3rd one shortly
           after the 2nd. I can't remember exactly why this caused a
           problem, but I'm hoping the sleep will take care of it.
-S         Pass -S to fai-chboot.
-w         Setup pxe, then wait like -a.
-h|--help  Print help and exit


Note, when switching between plain and arch, you will need to
do something like:
ssh wrt
cd /mnt/usb
rm tftpboot
ln -s <arch/debian iso dir> tftpboot


Note: Uses GNU getopt options parsing style
EOF
    exit $1
}

##### begin command line parsing ########

redep=true
acks=2
wait=false
temp=$(getopt -l help harSw "$@") || usage 1
eval set -- "$temp"
while true; do
    case $1 in
        -a) wait=true; set=false; shift ;;
        -r) redep=false; shift ;;
        -S) chboot_arg=-S; shift ;;
        -w) wait=true; set=true; shift ;;
        -h|--help) usage ;;
        --) shift; break ;;
        *) echo "$0: Internal error!" ; exit 1 ;;
    esac
done

read -r host type <<<"$@"

case $# in
    0|2);;
    *)
        echo "$0: error: expected 0 or 2 arguments"
        echo
        usage 1
        ;;
esac

if [[ $host && $host != default ]]; then
    host_tag="tag:$host,"
fi

##### end command line parsing ########

e() {
    echo "$@"
    "$@"
}

arch() {
    cat <<EOF
dhcp-option-force=209,boot/syslinux/archiso.cfg
dhcp-option-force=210,/arch/
dhcp-boot=${host_tag}/arch/boot/syslinux/lpxelinux.0
EOF
}

plain() {
    # if arch was used before, this additionally needs
    # the tftp link in /mnt/usb to be changed.
    cat <<EOF
dhcp-boot=${host_tag}pxelinux.0
EOF
}

fai() {
    cat <<EOF
dhcp-boot=${host_tag}pxelinux.0,faiserver.lan,faiserver.lan
EOF
}

ack-wait() {
    wait_count=$1
    if [[ $host ]]; then
        host_regex=" $host"
    fi
    regex=".*DHCPACK.*$host_regex$"
    i=0
    while (( i != wait_count )) && read -r line; do
        if [[ $line =~ $regex ]]; then
            i=$((i+1))
            echo $line
        fi
    done < <(ssh wrt logread -f)
    e sleep 20
}

set-pxe() {
    ${type:-:}|ssh wrt "cedit pxe /etc/dnsmasq.conf || /etc/init.d/dnsmasq restart
$([[ $type == arch ]] && echo arch-pxe-mount)"
}


if $set; then
    set-pxe
    if [[ $type == fai ]]; then
        e myfai-chboot $chboot_arg $host
        if $redep; then
            e fai-redep
        fi
    else
        # This will fail if faiserver is not setup, so ignore any
        # failure and don't bother us about it.
        myfai-chboot &>/dev/null ||:
    fi
fi

if $wait; then
    # fai's debian jessie 8.5ish does 2 dhcp requests when booting,
    # roughly 4 seconds apart. Earlier
    # versions did just 1. Now testing on a vm, it does 1.
    # bleh.
    echo "waiting for $acks dhcp acks then disabling pxe"
    ack-wait $acks
    type=
    set-pxe

    # previously tried waiting for one more ack then disabling faiserver,
    # since it can contain sensitive info, so turn it off when not in use,
    # but disabling that for now as it's inconvenient to clean this
    # up and run it in the background etc.

    # if [[ $type == fai ]]; then
    #     echo "waiting for 1 dhcp ack then disabling fai server"
    #     ack-wait 1
    #     faiserver-disable
    # fi
fi