#!/bin/sh
# copyright 2007 Canonical LTD., Oliver Grawert <ogra@ubuntu.com>
# distributed under the terms of the GNU General Public License version 2
# or any later version.

# Generates a squashfs image from an ltsp chroot to be served by an inetd driven
# nbd-server process.

usage() {
cat <<EOF
$0 [OPTION]
  -b, --basedir  Base of ltsp chroot.  Default is /opt/ltsp if unspecified.
  -p, --port     Port you wish this nbd image to communicate on.  Default
                 is 2000.
  -a, --arch     Architecture of this image.  Default is arch of the host.
  -h, --help     This message.
EOF
}

#
# Handle command line args
#

ARGS=$(getopt -o b:p:a:h --long base:,port:,arch:,help -n $0 -- "$@")

if [ $? != 0 ]; then
    exit 1                          # getopt failed
fi

eval set -- "${ARGS}"

while true ; do
    case "$1" in
        -b|--base) BASE=$2 ; shift 2 ;;
        -p|--port) PORT=$2 ; shift 2 ;;
        -a|--arch) ARCH=$2 ; shift 2 ;;
        -h|--help) usage ; exit 0 ;;
        --) shift ; break ;;
        *) echo "Internal error!" ; exit 1 ;;
    esac
done

# defaults
BASE=${BASE:-"/opt/ltsp"}

# make sure we dont carry trailing slashes around (LP 189237)
BASE=$(echo ${BASE}|sed -e 's/\/$//g')

PORT=${PORT:-"2000"}
if [ -z "${ARCH}" ]; then
    ARCH=$(dpkg --print-architecture)
fi

IMGDIR="${BASE}/images"
CONFFILE="/usr/share/initramfs-tools/conf.d/ltsp"
CHROOT="${BASE}/${ARCH}"

#
# for updateing the pxe config
#

TFTPBOOTDIR=${TFTPBOOTDIR:-"ltsp"}
TFTPDIRS=${TFTPDIRS:-"/var/lib/tftpboot /tftpboot"}

# source config file
if [ -f /etc/default/ltsp-update-image ]; then
    . /etc/default/ltsp-update-image
fi

if [ ! -d "${CHROOT}" ]; then
    echo "Error: chroot ${CHROOT} doesn't exist."
    exit 1
fi

# Find out if we are really an nbd client and if so, build the image.
# Check and see if user's left /proc mounted in chroot.  If so, issue
# a warning, and unmount it
PROC_MOUNTED=$(chroot /opt/ltsp/i386 test -f /proc/cpuinfo && echo True)
if [ -n "${PROC_MOUNTED}" ]; then
    echo "/proc mounted in chroot ${CHROOT}, Unmounting."
    umount ${CHROOT}/proc
fi

# Are we a chroot that wants an nbd image made?
if $(grep nbd "${CHROOT}/${CONFFILE}" > /dev/null 2>&1); then
    # make sure the images dir exists
    if [ ! -d ${IMGDIR} ]; then
        mkdir -p ${IMGDIR}
    fi

    IMAGE="${IMGDIR}/${ARCH}.img"
    rm "${IMAGE}"
    mksquashfs "${CHROOT}" "${IMAGE}.tmp" -e cdrom
    mv "${IMAGE}.tmp" "${IMAGE}"
    chmod 0744 "${IMAGE}"

    # Make sure hosts.allow has the keepalive option for nbdrootd
    if [ -z "$(grep nbdrootd /etc/hosts.allow)" ]; then
        echo 'nbdrootd: ALL: keepalive' >> /etc/hosts.allow
    fi

    # make sure image and port are properly defined in inetd.conf
    if grep ^${PORT} /etc/inetd.conf |grep ${IMAGE} > /dev/null 2>&1 ; then
        echo "Info: port ${PORT} is already defined with ${IMAGE} in inetd.conf"
        echo "Info: taking no action."
    else
        if [ -n "$(grep ${IMAGE} /etc/inetd.conf | grep -v '^#')" ] ; then
            # if its already there, use its settings
            echo "Warning: image defined with different port already in inetd.conf"
            PORT=$(grep ${IMAGE} /etc/inetd.conf |cut -d' ' -f1)
            echo "Info: using already defined portnumber ${PORT}"
        else
            if grep ^${PORT} /etc/inetd.conf > /dev/null 2>&1 ; then
                # make sure we use a free port
                echo "Warning: found port in use already"
                while grep ^${PORT} /etc/inetd.conf > /dev/null 2>&1; do
                    PORT=$(expr ${PORT} + 1)
                    echo "Info: raising nbdroot portnumber by one to be ${PORT}"
                done
            fi
            echo "Info: updating inetd config"
            update-inetd --group LTSP --add "${PORT}               stream  tcp\
            nowait  nobody /usr/sbin/tcpd /usr/sbin/nbdrootd ${IMAGE}"
        fi
    fi

    if [ ${PORT} -ne 2000 ];then
        # seems we use a non default port, add it to pxe config and chroot
        echo "Info: updating PXE"
        for TFTPDIR in $TFTPDIRS ; do
            if [ ! -d $TFTPDIR ]; then
                continue       # skip directory
            fi
            PXECFG="${TFTPDIR}/${TFTPBOOTDIR}/${ARCH}/pxelinux.cfg/default"
            if [ ! -f ${PXECFG} ]; then
                continue
            fi
            if grep nbdport ${PXECFG} > /dev/null 2>&1 ; then
                sed -i -e "s/nbdport=[0-9]*/nbdport=${PORT}/g" ${PXECFG}
            else
                sed -i -e "s/$/ nbdport=${PORT}/g" ${PXECFG}
            fi
        done
        if [ ! -d "${CHROOT}/etc/ltsp" ]; then
            mkdir "${CHROOT}/etc/ltsp"
        fi
        echo "BOOTPROMPT_OPTS=\"quiet splash nbdport=${PORT}\"" >${CHROOT}/etc/ltsp/update-kernels.conf
    fi
fi
