#!/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, fresh pxelinux.cfg files and restart nbd-server.
  -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.
  -N, --nbd-named     Use nbd-server named entries
  -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:nNhifo:O:S:T: --long base:,port:,arch:,no-comp,tftpbootdir:,exclude-dirs:,force,image-only,options:,ipappend:,no-comp,nbd-named,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 ;;
        -N|--nbd-named) NBD_NAMED="1" ; 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"}
NBD_NAMED=${NBD_NAMED:-"0"}

# 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}"

# Generate nbd-server configuration
if [ "$NBD_NAMED" = "1" ] && [ ! -d /etc/nbd-server ]; then
    echo "Your system doesn't have a /etc/nbd-server directory."
    echo "Assuming inetd is needed to export the image"
    NBD_NAMED="0"
else
    NBD_NAME="ltsp_$ARCH"
fi

#
# 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
}

assign_name () {
    THIS_CHROOT=$CHROOT

    if [ ! -d "${THIS_CHROOT}/etc/ltsp" ]; then
        mkdir "${THIS_CHROOT}/etc/ltsp"
    fi

    if [ -n "${NBD_SERVER}" ]; then
        echo "BOOTPROMPT_OPTS=\"${BOOTPROMPT_OPTIONS} nbdserver=${NBD_SERVER} nbdname=${NBD_NAME} \"" >${THIS_CHROOT}/etc/ltsp/update-kernels.conf
    else
        echo "BOOTPROMPT_OPTS=\"${BOOTPROMPT_OPTIONS} nbdname=${NBD_NAME} \"" >${THIS_CHROOT}/etc/ltsp/update-kernels.conf
    fi

    echo "NBD_NAME=${NBD_NAME}" >>${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 "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 nbdname ${PXECFG} > /dev/null 2>&1 ; then
            sed -i -e "s/nbdname=.* /nbdname=${NBD_NAME} /g" ${PXECFG}
        fi
    done
    echo "Done."
    RUN_LUK=1
}

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

        if [ "$NBD_NAMED" = "1" ]; then
            /usr/share/ltsp/scripts/nbd-update-config --add $NBD_NAME $IMAGE
            echo "The nbd-server daemon doesn't allow for live configuration update."
            echo "To pick up the changes, the daemon needs to be restarted."
            echo "If it's a new LTSP chroot. Run /etc/init.d/nbd-server restart"
            echo "THIS WILL DISCONNECT ALL YOUR CLIENTS (they'll need to be rebooted)"
        fi
    fi
}

if [ -z "${IMAGE_ONLY}" ]; then
    if [ "$NBD_NAMED" = "0" ]; then
        collect_ports
        assign_ports
    else
        assign_name
    fi
    [ -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

