#!/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>
#  2012, Alkis Georgopoulos <alkisg@gmail.com>
#
#  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 configuration 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 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.
  -h, --help          This message.
  -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.
  -O, --openbsd-inetd Use old inetd method instead of named NBD exports.
  -p, --port          Port you wish this nbd image to communicate on.  Default
                      is 2000.
  -s, --skip-image    Do not generate an image, but perform other actions.
  -S, --server        Specify the NBD server IP.
  -T, --timeout       Add timeout to bootloader.
EOF
}

#
# Handle command line args
#

ARGS=$(getopt -n "$0" -o a:b:d:e:fhiI:no:Op:sS:T:h --long \
    arch:,base:,tftpbootdir:,exclude-dirs:,force,help,image-only,ipappend:,no-comp,options:,openbsd-inetd,port:,skip-image,server:,timeout: -- "$@")

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

eval set -- "${ARGS}"

while true ; do
    case "$1" in
        -a|--arch) ARCH=$2 ; shift 2 ;;
        -b|--base) BASE=$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 ;;
        -h|--help) usage ; exit 0 ;;
        -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 ;;
        -O|--openbsd-inetd) NBD_NAMED="0" ; shift 1 ;;
        -p|--port) PORT=$2 ; shift 2 ;;
        -s|--skip-image) SKIP_IMAGE=True ; shift 1 ;;
        -S|--server)  NBD_ROOT_HOST=$2 ; shift 2 ;;
        -T|--timeout) TIMEOUT=$2 ; shift 2 ;;
        --) shift ; break ;;
        *) echo "Internal error!" ; exit 1 ;;
    esac
done

if [ "$LTSP_UPDATE_IMAGE_DISABLED" = "True" ]; then
    echo "$0 is currently disabled. To enable, edit /etc/ltsp/ltsp-update-image.conf."
    exit 0
fi

# defaults
BASE=${BASE:-"/opt/ltsp"}
BOOTPROMPT_OPTIONS=${BOOTPROMPT_OPTIONS:-"root=/dev/nbd0 init=/sbin/init-ltsp quiet splash plymouth:force-splash vt.handoff=7"}
NBD_NAMED=${NBD_NAMED:-"1"}

# 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"
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_ROOT_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_ROOT_PORT will be recorded by this script in the chroot
        # For future reference
        unset NBD_ROOT_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_ROOT_PORT is a number
        NBD_ROOT_PORT=$(echo ${NBD_ROOT_PORT}|sed -e '/[^0-9]/d')
        ALL_PORT_CHROOTS="${ALL_PORT_CHROOTS} ${NBD_ROOT_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_ROOT_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_ROOT_PORT=${c%%=*}
        THIS_CHROOT=${c#*=}
        case "${PORTS_FOUND}" in
            *-${NBD_ROOT_PORT}-*) 
                OLD_PORT=${NBD_ROOT_PORT}
                find_available_port
                if [ -n "${OLD_PORT}" ]; then
                    echo "Port ${OLD_PORT} already in use. Changing to port ${NBD_ROOT_PORT}."
                else
                    echo "Cannot determine assigned port. Assigning to port ${NBD_ROOT_PORT}."
                fi

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

                if [ -n "${NBD_ROOT_HOST}" ]; then
                    echo "BOOTPROMPT_OPTS=\"${BOOTPROMPT_OPTIONS} nbdroot=${NBD_ROOT_HOST}:${NBD_ROOT_PORT}\"" >${THIS_CHROOT}/etc/ltsp/update-kernels.conf
                else
                    echo "BOOTPROMPT_OPTS=\"${BOOTPROMPT_OPTIONS} nbdport=${NBD_ROOT_PORT}\"" >${THIS_CHROOT}/etc/ltsp/update-kernels.conf
                fi

                echo "NBD_ROOT_PORT=${NBD_ROOT_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."
                # run ltsp-update-kernels later...
                RUN_LUK=1

                echo -n "Configuring inetd... "
                IMAGE=${IMGDIR}/${THIS_CHROOT##*/}.img 
                sed -i -e "\|${IMAGE}|d" /etc/inetd.conf
                update-inetd --group LTSP --add "${NBD_ROOT_PORT}               stream  tcp nowait  nobody /usr/sbin/tcpd /usr/sbin/nbdrootd ${IMAGE}"
                echo "Done."

                ;;
        esac
        PORTS_FOUND="${PORTS_FOUND}${NBD_ROOT_PORT}-"
    done
}

assign_name () {
    THIS_CHROOT=$CHROOT

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

    if [ -n "${NBD_ROOT_HOST}" ]; then
        echo "BOOTPROMPT_OPTS=\"${BOOTPROMPT_OPTIONS} nbdroot=${NBD_ROOT_HOST}:${NBD_ROOT_NAME} \"" >${THIS_CHROOT}/etc/ltsp/update-kernels.conf
    else
        echo "BOOTPROMPT_OPTS=\"${BOOTPROMPT_OPTIONS} nbdroot=:${NBD_ROOT_NAME} \"" >${THIS_CHROOT}/etc/ltsp/update-kernels.conf
    fi

    echo "NBD_ROOT_NAME=${NBD_ROOT_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."
    RUN_LUK=1
}

# Creates an nbd-server configuration file for the specified nbd image, if it
# doesn't already exist. If it does, no action is performed.
# Additionally, if /etc/nbd-server/config doesn't exist, it creates it,
# it generates a configuration file for swap, and it starts nbd-server."
create_nbd_config() {
    local exportname section conffile

    exportname="$1"
    test -f "$exportname" || return 1
    section=${exportname##*/}
    section=${section%.*}
    conffile="/etc/nbd-server/conf.d/ltsp_$section.conf"
    if [ -f "$conffile" ]; then
        echo "Configuration file $conffile already exists, no action taken."
        return
    fi

    cat >"$conffile" <<EOF
[ltsp_$section]
exportname = $exportname
readonly = true
EOF
    echo "A new configuration file has been created for image $exportname."

    if [ -f "/etc/nbd-server/config" ]; then
        echo \
"For nbd-server to re-read its configuration, you need to manually run:
    service nbd-server restart
THIS WILL DISCONNECT ALL RUNNING CLIENTS (they'll need to be rebooted)."
        return
    fi
    
    # if user/group nbd exist, configure 
    nbd_user=$(getent passwd nbd | cut -d : -f 1)
    nbd_group=$(getent group nbd | cut -d : -f 1)
    test -z "$nbd_user" && nbd_user=nobody
    test -z "$nbd_group" && nbd_group=nogroup

    cat > "/etc/nbd-server/config" <<EOF
[generic]
user = $nbd_user
group = $nbd_group
oldstyle = false
includedir = /etc/nbd-server/conf.d
EOF

    if [ ! -f "/etc/nbd-server/conf.d/swap.conf" ]; then
        mkdir -p "/etc/nbd-server/conf.d"
        cat > "/etc/nbd-server/conf.d/swap.conf" <<EOF
[swap]
exportname = /tmp/nbd-swap/%s
prerun = nbdswapd %s
postrun = rm -f %s
EOF
    fi

    # We can assume that nbd-server wasn't running without a config file
    if ! invoke-rc.d nbd-server start; then
        echo "Failed to start nbd-server." >&2
    fi
}

generate_image () {
    # Check if the chroot initramfs supports nbd booting
    test -f "${CHROOT}/usr/share/initramfs-tools/scripts/local-top/nbd" || return

    # 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

    # 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
    test -x /usr/bin/nice && nice=nice || unset nice
    test -x /usr/bin/ionice && ionice=ionice || unset ionice
    $nice $ionice mksquashfs "${CHROOT}" "${IMAGE}.tmp" $NO_COMP -e cdrom ${EX_DIRS}
    if [ "$?" != "0" ]; then
        rm -f "${IMAGE}.tmp"
        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
        create_nbd_config "$IMAGE"
    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

if [ "$SKIP_IMAGE" != "True" ]; then
    generate_image
fi

# 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

