#!/bin/sh
#
#  Copyright (c) 2007 Canonical LTD
#
#  Author: Oliver Grawert <ogra@canonical.com>
#
#  2007, Scott Balneaves <sbalneav@ltsp.org>
#        Warren Togami <wtogami@redhat.com>
#  2008, Vagrant Cascadian <vagrant@freegeek.org>
#  2010, Gideon Romm <gadi@ltsp.org>
#
#  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, you can find it on the World Wide
#  Web at http://www.gnu.org/copyleft/gpl.html, or write to the Free
#  Software Foundation, Inc., 51 Franklin St, Fifth Floor, Boston,
#  MA 02110-1301, USA.
#

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

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



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

usage() {
cat <<EOF
$0 [OPTION]
  -a, --arch          Architecture of this image.  Default is arch of the host.
  -b, --base          Base of ltsp chroot.  Default is /opt/ltsp if unspecified.
  -d, --tftpbootdir   Subdir within tftpdir where ltsp kernels are.  Defaults
                      to "ltsp".
  -e, --exclude-dirs  Exclude those dirs from the image.
  -f, --force         Force recalculating ports for all chroots, update of all 
                      kernels and fresh pxelinux.cfg files.
  -i, --image-only    Update image only and do not attempt to adjust ports or 
                      services.
  -I, --ipappend      Pass IPAPPEND value to bootloader.
  -n, --no-comp       Do not compress the image.
  -o, --options       Pass kernel command line options to all config files.
  -p, --port          Port you wish this nbd image to communicate on.  Default
                      is 2000.
  -S, --server        Specify the NBD server IP.
  -T, --timeout       Add timeout to bootloader.
  -h, --help          This message.
EOF
}

#
# Handle command line args
#

ARGS=$(getopt -o b:p:a:e:nhifo:O:S:T: --long base:,port:,arch:,no-comp,tftpbootdir:,exclude-dirs:,force,image-only,options:,ipappend:,server:,timeout:,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 ;;
        -d|--tftpbootdir) TFTPBOOTDIR=$2 ; shift 2 ;;
        -e|--exclude-dirs) EX_DIRS=$(echo $2| sed -e 's/dev\|tmp\|proc\|var//g') ; shift 2 ;;
        -f|--force) FORCE_NEW=1; shift 1 ;;
        -i|--image-only) IMAGE_ONLY=1; shift 1 ;;
        -I|--ipappend) IPAPPEND=$2 ; shift 2 ;;
        -n|--no-comp) NO_COMP="-noF -noD -noI -no-exports" ; shift 1 ;;
        -o|--options) BOOTPROMPT_OPTIONS=$2 ; shift 2 ;;
        -S|--server)  NBD_SERVER=$2 ; shift 2 ;;
        -T|--timeout) TIMEOUT=$2 ; shift 2 ;;
        -h|--help) usage ; exit 0 ;;
        --) shift ; break ;;
        *) echo "Internal error!" ; exit 1 ;;
    esac
done

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

# 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 updating the pxe config
#

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

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

collect_ports () {
    ALL_CHROOTS=${ALL_CHROOTS:-"$(find $BASE/ -mindepth 1 -maxdepth 1 -type d | grep -v images)"}
    unset ALL_PORT_CHROOTS
    for c in ${ALL_CHROOTS}; do
        # NBD_PORT will be recorded by this script in the chroot
        # For future reference
        unset NBD_PORT 
        [ -n "${FORCE_NEW}" ] && rm -f ${c}/etc/ltsp/update-kernels.conf 2>/dev/null
        [ -f ${c}/etc/ltsp/update-kernels.conf ] && . ${c}/etc/ltsp/update-kernels.conf
        # Make sure NBD_PORT is a number
        NBD_PORT=$(echo ${NBD_PORT}|sed -e '/[^0-9]/d')
        ALL_PORT_CHROOTS="${ALL_PORT_CHROOTS} ${NBD_PORT}=${c}"
    done 
}

find_available_port () {
    unset FOUND_PORT
    # Set p to the base port
    p=${PORT}
    while [ -z "${FOUND_PORT}" ]; do
        case "${PORTS_FOUND}" in
            *-${p}-*) p=$(($p+1)) ;;
            *) 
                if netstat -lnp|grep :${p} >/dev/null && [ -z "$(grep ${p} /etc/inetd.conf| grep nbdrootd)" ]; then 
                    # Port is currently in use by another service
                    p=$(($p+1)); continue
                else 
                    export NBD_PORT=$p; FOUND_PORT=1; return 0 
                fi
                ;;
        esac
    done
}

assign_ports () {
    PORTS_FOUND="--"
    unset RUN_LUK
    for c in $(echo ${ALL_PORT_CHROOTS}|tr ' ' '\n'| sort); do
        NBD_PORT=${c%%=*}
        THIS_CHROOT=${c#*=}
        case "${PORTS_FOUND}" in
            *-${NBD_PORT}-*) 
                OLD_PORT=${NBD_PORT}
                find_available_port
                if [ -n "${OLD_PORT}" ]; then
                    echo "Port ${OLD_PORT} already in use. Changing to port ${NBD_PORT}."
                else
                    echo "Cannot determine assigned port. Assigning to port ${NBD_PORT}."
                fi
                if [ ! -d "${THIS_CHROOT}/etc/ltsp" ]; then
                    mkdir "${THIS_CHROOT}/etc/ltsp"
                fi
                if [ -n "${NBD_SERVER}" ]; then
                    echo "BOOTPROMPT_OPTS=\"${BOOTPROMPT_OPTIONS} nbdroot=${NBD_SERVER}:${NBD_PORT}\"" >${THIS_CHROOT}/etc/ltsp/update-kernels.conf
                else
                    echo "BOOTPROMPT_OPTS=\"${BOOTPROMPT_OPTIONS} nbdport=${NBD_PORT}\"" >${THIS_CHROOT}/etc/ltsp/update-kernels.conf
                fi
                echo "NBD_PORT=${NBD_PORT}" >>${THIS_CHROOT}/etc/ltsp/update-kernels.conf
                [ -n "${TIMEOUT}" ] && echo "TIMEOUT=${TIMEOUT}" >>${THIS_CHROOT}/etc/ltsp/update-kernels.conf
                [ -n "${IPAPPEND}" ] && echo "IPAPPEND=${IPAPPEND}" >>${THIS_CHROOT}/etc/ltsp/update-kernels.conf
                echo "Regenerating kernel... "
                chroot ${THIS_CHROOT} /usr/share/ltsp/update-kernels
                echo "Done."
                echo -n "Configuring inetd... "
                IMAGE=${IMGDIR}/${THIS_CHROOT##*/}.img 
                sed -i -e "\|${IMAGE}|d" /etc/inetd.conf
                update-inetd --group LTSP --add "${NBD_PORT}               stream  tcp nowait  nobody /usr/sbin/tcpd /usr/sbin/nbdrootd ${IMAGE}"
                echo "Done."
                echo -n "Updating pxelinux default configuration..."
                for TFTPDIR in $TFTPDIRS ; do
                    if [ ! -d $TFTPDIR ]; then
                        continue       # skip directory
                    fi
                    PXECFG="${TFTPDIR}/${TFTPBOOTDIR}/${THIS_CHROOT##*/}/pxelinux.cfg/default"
                    [ -n "${FORCE_NEW}" ] && rm -f ${PXECFG} 2>/dev/null
                    if [ ! -f ${PXECFG} ]; then
                        mkdir -p "${TFTPDIR}/${TFTPBOOTDIR}/${THIS_CHROOT##*/}/pxelinux.cfg"
                        [ -n "${TIMEOUT}" ] && TIMEOUT_LINE="timeout ${TIMEOUT}"
                        [ -n "${IPAPPEND}" ] && IPAPPEND_LINE="IPAPPEND ${IPAPPEND}"
                        echo <<EOF >${PXECFG}
default ltsp 
${TIMEOUT_LINE}

label ltsp 
kernel vmlinuz 
append ro initrd=initrd.img ${BOOTPROMPT_OPTS}
${IPAPPEND_LINE}
EOF
                        continue
                    fi
                    if grep nbdport ${PXECFG} > /dev/null 2>&1 ; then
                        sed -i -e "s/nbdport=[0-9]*/nbdport=${NBD_PORT}/g" ${PXECFG}
                    fi
                    if grep nbdroot ${PXECFG} > /dev/null 2>&1 ; then
                        sed -i -e "s/nbdroot=\([^:]\+\):[0-9]*/nbdroot=\1:${NBD_PORT}/" ${PXECFG}
                    fi
                done
                echo "Done."
                RUN_LUK=1
                ;;
        esac
        PORTS_FOUND="${PORTS_FOUND}${NBD_PORT}-"
    done 
}

generate_image () {
    # 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 ${CHROOT} test -f /proc/cpuinfo && echo True)
    if [ -n "${PROC_MOUNTED}" ]; then
        echo "/proc mounted in chroot ${CHROOT}, unmounting."
        # binfmt_misc might need to be unmounted manually, see LP #534211
        if grep -q "^binfmt_misc $ROOT/proc/sys/fs/binfmt_misc" $ROOT/proc/mounts; then
            umount $ROOT/proc/sys/fs/binfmt_misc || true
        fi
        if ! umount ${CHROOT}/proc; then
            echo "${CHROOT}/proc is in use, forcing unmount." >&2
            umount -l ${CHROOT}/proc
        fi
    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 -f "${IMAGE}.tmp" >/dev/null 2>&1
        nice mksquashfs "${CHROOT}" "${IMAGE}.tmp" $NO_COMP -e cdrom ${EX_DIRS}
        if [ "$?" != "0" ]; then
            echo "Error: mksquashfs failed to build the ltsp image, exiting"
            exit 1
        fi
        if [ -f "${IMAGE}.tmp" ]; then
            mv "${IMAGE}.tmp" "${IMAGE}"
            chmod 0644 "${IMAGE}"
        fi
    fi
}

if [ -z "${IMAGE_ONLY}" ]; then
    collect_ports
    assign_ports
    [ -n "$RUN_LUK" ] && ltsp-update-kernels -b "${BASE}" -t "${TFTPDIRS}" -d "${TFTPBOOTDIR}"
fi

generate_image

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

