#!/bin/sh
#
#  Copyright (c) 2005 Canonical LTD
#
#  Author: Matt Zimmerman <mdz@canonical.com>
#
#  2006, Oliver Grawert <ogra@canonical.com>
#        Vagrant Cascadian <vagrant@freegeek.org>
#  2007, Scott Balneaves <sbalneav@ltsp.org>
#        Oliver Grawert <ogra@canonical.com>
#  2008, Vagrant Cascadian <vagrant@freegeek.org>
#        Warren Togami <wtogami@redhat.com>
#        Oliver Grawert <ogra@canonical.com>
#  2009, Warren Togami <wtogami@redhat.com>
#  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.
#

usage() {
    cat <<EOF
Usage: $0 [OPTION] [CHROOT...]

Copies the boot/ directory from LTSP chroots to the TFTP directories
in order to make them available to PXE clients. Copying kernels from
inside NBD images is also supported. CHROOT can be a fullpath or a
subdirectory of the base directory, and if it's unset, all available
chroots are processed.

Options:
  -b, --base[=PATH]     The LTSP base directory. Defaults to /opt/ltsp if unspecified.
  -h, --help            Displays the ltsp-update-kernels help message.
      --version         Output version information and exit.
EOF
}

trap_cleanup() {
    # Stop trapping
    trap - 0 HUP INT QUIT KILL SEGV PIPE TERM
    umount_marked
    rmdir "$mnt"
}

update_kernels() {
    local name tftpdir tftpboot chroot 
    name=$1
    tftpdir=$2

    tftpboot="$tftpdir/$TFTP_BOOT_DIR"
    tftpboot=${tftpboot%/}

    # Loop-mounting NBD images in order to copy their kernels is supported
    # in order to allow for btrfs or ext loopback images, or images transferred
    # from other sources to the LTSP server.
    # But if both the chroot and the NBD image exist, the chroot is preferred,
    # to make updating BOOTPROMPT_OPTS in update-kernels.conf easier.
    # A PREFER_NBD_IMAGE environment variable is supported though, to make it
    # possible for `ltsp-update-image --revert` to use the NBD kernels.
    unset chroot
    if [ -x "$BASE/$name/bin/true" ] && [ "$PREFER_NBD_IMAGE" != true ]; then
        chroot="$BASE/$name"
    elif [ -f "$BASE/images/$name.img" ]; then
        if [ -z "$mnt" ]; then
            mnt=$(mktemp -d)
            trap "trap_cleanup" 0 HUP INT QUIT KILL SEGV PIPE TERM
        else
            umount_marked
        fi
        if mark_mount -o loop,ro "$BASE/images/$name.img" "$mnt"; then
            chroot="$mnt"
        fi
    fi
    if [ -z "$chroot" ]; then
        echo "Skipping invalid chroot: $name"
        continue
    fi
    echo "Updating $tftpdir directories for chroot: $name"

    # Source distro-specific variables from the chroot
    unset KERNEL_NAMES INITRD_NAME
    if [ -f "$chroot/etc/ltsp/update-kernels.conf" ]; then
        . "$chroot/etc/ltsp/update-kernels.conf"
    fi
    mkdir -p "$tftpboot/$name"
    cp -a "$chroot/boot/." "$tftpboot/$name/"

    # Generate pxelinux.cfg/default symlink if not present.
    if [ -f "$tftpboot/$name/pxelinux.cfg/ltsp" ]; then
        # Remove autogenerated pxelinux.cfg/default
        if [ -f "$tftpboot/$name/pxelinux.cfg/default" ]; then
            if [ ! -L "$tftpboot/$name/pxelinux.cfg/default" ]; then
                if grep -q '# This file is regenerated when update-kernels runs.' "$tftpboot/$name/pxelinux.cfg/default" ; then
                    rm -f "$tftpboot/$name/pxelinux.cfg/default"
                fi
            fi
        fi
        if [ ! -f "$tftpboot/$name/pxelinux.cfg/default" ]; then
            ln -sf ltsp "$tftpboot/$name/pxelinux.cfg/default" 
        fi
    fi

    # Ensure that the files are readable (LP: #759115) (Dracut initramfs)
    find "$tftpboot/$name/" -maxdepth 1 ! -perm -o=r -exec chmod a+r {} \;

    # OFW on Mac is lame, they cannot tftp from directories
    if [ -e "$tftpboot/$name/yaboot" ]; then
        if [ ! -e "$tftpdir/yaboot" ]; then
            ln -sf "$TFTP_BOOT_DIR/$name/yaboot" "$tftpdir/yaboot"
        fi
        if [ ! -e "$tftpdir/yaboot.conf" ]; then
            ln -sf "$TFTP_BOOT_DIR/$name/yaboot.conf" "$tftpdir/yaboot.conf"
        fi
    fi

    # Cleanup old kernels and images from tftpboot directory
    cleanup_kernels "$name" "$tftpboot" "$chroot"

    link_kernel_flavors "$tftpboot/$name"
}

# Create symlinks for each kernel flavor in the tftp dir.
# It requires that the distro-specific $KERNEL_NAMES and $INITRD_NAME variables
# are declared in $CHROOT/etc/ltsp/update-kernels.conf.
link_kernel_flavors() {
    local tftpname last_flavor file name version flavor initrd

    tftpname=$1
    if [ ! -d "$tftpname" ]; then
        echo "Directory $tftpname does not exist"
        return 1
    fi

    # Those defaults should work on debian-based distros, but shouldn't hurt
    # elsewhere because they wouldn't match actual files and they'd be ignored.
    KERNEL_NAMES=${KERNEL_NAMES:-'s/\(vmlinu[xz]-\)\([^-]*-[^-]*-\)\(.*\)/& \1 \2 \3/p'}
    INITRD_NAME=${INITRD_NAME:-'s/vmlinu[xz]/initrd.img/p'}
    # 4.4.0-159 to 4.4.0-166 kernels have a tmpfs/overlay bug
    BANNED_NAMES="4.4.0-1(59|6[1456])"


    last_flavor=
    find "$tftpname" -mindepth 1 -maxdepth 1 -type f -printf "%f\n" \
    | sed -n "$KERNEL_NAMES" | grep -v -E "$BANNED_NAMES" | sort -k 4,4V -k 3,3rV \
    | while read file name version flavor; do
        if [ "$flavor" != "$last_flavor" ]; then
            initrd=$(echo "$file" | sed -n "$INITRD_NAME")
            if [ ! -e "$tftpname/$initrd" ]; then
                echo "Ignoring $file because the matching $initrd doesn't exist" >&2
                continue
            fi
            ln -sf "$file" "$tftpname/$name$flavor"
            ln -sf "$initrd" "$tftpname/$(echo "$name$flavor" | sed -n "$INITRD_NAME")"
            last_flavor=$flavor
        fi
    done
}

# distro specific functions

# For all kernels in TFTP, find the kernel $version from vmlinuz-* filename.
# If the corresponding /opt/ltsp/$name/lib/modules/$version is missing,
# then delete kernel and images for this version from tftpboot directory.
# Distros that don't match vmlinuz-* should override this function.
cleanup_kernels() {
    local name tftpboot chroot
    name=$1
    tftpboot=$2
    chroot=$3

    # Loop through every vmlinuz-* file
    for kernelpath in $(find "$tftpboot/$name/" -type f -name 'vmlinuz-*' -o -name 'vmlinux-*' | grep -v -E "$BANNED_NAMES"); do
        kernel=${kernelpath##*/}
        case $kernel in
            vmlinuz-*) version=${kernel#vmlinuz-} ;;
            vmlinux-*) version=${kernel#vmlinux-} ;;
        esac
        if [ ! -d "$chroot/lib/modules/$version" ]; then
            echo "Removing $kernelpath"
            # Common
            rm -f "$tftpboot/$name/$kernel"
            rm -f "$tftpboot/$name/config-$version"
            rm -f "$tftpboot/$name/System.map-$version"
            rm -rf "$tftpboot/$name/dtbs-$version/"
            # Fedora
            rm -f "$tftpboot/$name/initrd-$version.img"
            rm -f "$tftpboot/$name/initramfs-$version.img"
            rm -f "$tftpboot/$name/elf-$version.img"
            rm -f "$tftpboot/$name/wraplinux-nbi-$version.img"
            rm -f "$tftpboot/$name/aout-$version.img"
            rm -f "$tftpboot/$name/symvers-$version.gz"
            # Debian
            rm -f "$tftpboot/$name/initrd.img-$version"
            rm -f "$tftpboot/$name/nbi.img-$version"
            # Ubuntu
            rm -f "$tftpboot/$name/abi-$version"
            rm -f "$tftpboot/$name/vmcoreinfo-$version"
        fi
    done
}

# Set an optional MODULES_BASE, so help2man can be called from build env
MODULES_BASE=${MODULES_BASE:-/usr/share/ltsp}

# This also sources vendor functions and .conf file settings
. ${MODULES_BASE}/ltsp-server-functions

if ! args=$(getopt -n "$0" -o b:h \
    -l "base:,help,version" -- "$@")
then
    exit 1
fi
eval "set -- $args"
while true ; do
    case "$1" in
        -b|--base) shift; BASE=$1 ;;
        -h|--help) usage; exit 0 ;;
        --version) ltsp_version; exit 0 ;;
        --) shift ; break ;;
        *) die "$0: Internal error!" ;;
    esac
    shift
done
require_root

BASE=${BASE:-"/opt/ltsp"}
# Remove trailing /, if present
BASE=${BASE%/}

# Chroots can be specified in the command line. If not, update all of them.
if [ $# -eq 0 ]; then
    set -- $(
        {
            find -L "$BASE/" -mindepth 1 -maxdepth 1 -type d ! -name images \
                ! -name lost+found -printf "%f\n"
            if [ -d "$BASE/images/" ]; then
                find -L "$BASE/images/" -mindepth 1 -maxdepth 1 -type f \
                    -name '*.img' -printf "%f\n" | sed 's/.img$//'
            fi
        } | sort -u
    )
fi
test $# -gt 0 || die "No chroots found in $BASE"

for tftpdir in $TFTP_DIRS; do
    if [ ! -d "$tftpdir" ]; then
        # skip directory
        continue
    fi
    for name in "$@"; do
        update_kernels "$name" "$tftpdir"
    done
    # Update selinux file contexts if necessary
    if [ -f /selinux/enforce ] && [ -x /sbin/restorecon ]; then
        restorecon -R "$tftpdir" > /dev/null
    fi
done