#!/usr/bin/env bash
# Copyright 2012, 2013 Canonical Ltd.  This software is licensed under the
# GNU Affero General Public License version 3 (see the file LICENSE).
#
# Download static files needed for net-booting nodes through TFTP:
# pre-boot loader, kernels, and initrd images.
#
# This script downloads the required files into the TFTP home directory
# (by default, /var/lib/maas/tftp).  Run it with the necessarily privileges
# to write them there.

# Exit immediately if a command exits with a non-zero status.
set -o errexit
# Treat unset variables as an error when substituting.
set -o nounset

# Load settings if available.
settings="/etc/maas/import_pxe_files"
[ -r $settings ] && . $settings

# Location of the GPG keyring for the Ubuntu archive.
GPG_KEYRING="${GPG_KEYRING:-/usr/share/keyrings/ubuntu-archive-keyring.gpg}"

# Whether to skip checking GPG keys (necessary for testing of this script).
# Set this to a nonempty string to skip checking.  This is needed for test
# programs, because they can't possibly sign with the right key.
IGNORE_GPG="${IGNORE_GPG:-}"

# Download locations for Ubuntu releases.  When the cluster controller runs
# the import scripts, it provides settings from the server side.
MAIN_ARCHIVE=${MAIN_ARCHIVE:-http://archive.ubuntu.com/ubuntu/}
PORTS_ARCHIVE=${PORTS_ARCHIVE:-http://ports.ubuntu.com/ubuntu-ports/}

# Ubuntu releases that are to be downloaded.
SUPPORTED_RELEASES=$(distro-info --supported)
RELEASES=${RELEASES:-$SUPPORTED_RELEASES}

# The current Ubuntu release.
STABLE_RELEASE=$(distro-info --stable)

# Supported architectures.
# armhf/highbank is also supported by this script, but cannot be enabled here
# until maas-import-ephemerals also supports it or IMPORT_EPHEMERALS is set to
# 0.
ARCHES=${ARCHES:-amd64/generic i386/generic armhf/highbank}

# Command line to download a resource at a given URL into the current
# directory.  A wget command line will work here, but curl will do as well.
DOWNLOAD=${DOWNLOAD:-wget --no-verbose}

# Whether to download ephemeral images as well: "1" for yes, "0" for no.
# Default is yes.
IMPORT_EPHEMERALS=${IMPORT_EPHEMERALS:-1}

fail() {
    local msg=$1

    echo $msg >&2
    exit 1
}

# Show script usage/summary.
show_usage() {
    echo "Usage: ${0##*/}"
    echo
    echo "This helper script downloads the relevant boot images from an "
    echo "Ubuntu archive and uses 'maas' to provision them for PXE booting "
    echo "from TFTP."
    echo
    echo "This script takes no arguments, but you can adjust some parameters "
    echo -e "by editing the config file found at \033[1m$settings\033[0m."
    echo
    echo "MAAS homepage:<http://maas.ubuntu.com>"
    echo
}

# Return a URL that points to the images directory for the relevant
# release/arch.
compose_installer_base_url() {
    local arch=$1 release=$2

    case $arch in
        amd64/*|i386/*)
            local installer_url="$MAIN_ARCHIVE/dists/$release/main/installer-${arch%%/*}"
            echo "$installer_url/current/images/"
            ;;
        armhf/*)
            # No ARM server installers were available in precise, so always go for -updates for now
            # A better general fix is LP: #1052397
            if [ "$release" = "precise" ]; then
                updates=-updates
            else
                updates=
            fi
            local installer_url="$PORTS_ARCHIVE/dists/${release}${updates}/main/installer-${arch%%/*}"
            echo "$installer_url/current/images/"
            ;;
        *)
            echo "Unknown architecture: $arch" >&2
            exit 1
            ;;
    esac
}

# Return the URL part that is appended to the base url that gives the location
# of the images.
compose_installer_download_url_postfix() {
    local arch=$1

    case $arch in
        amd64/*|i386/*)
            echo "netboot/ubuntu-installer/${arch%%/*}/"
            ;;
        armhf/*)
            echo "${arch#*/}/netboot/"
            ;;
        *)
            echo "Unknown architecture: $arch" >&2
            exit 1
            ;;
    esac
}

# Put together a full URL for where the installer files for architecture $1
# and release $2 can be downloaded.
compose_installer_download_url() {
    local arch=$1 release=$2

    base_url=$(compose_installer_base_url $arch $release)
    postfix=$(compose_installer_download_url_postfix $arch)

    echo "$base_url/$postfix"
}

# Fetch MD5SUMS file.  This returns false (i.e. nonzero) in the case of
# survivable failure, so that the caller can skip this image and move on.
fetch_server_md5sums() {
    local base_url=$1

    if ! $DOWNLOAD "$base_url/MD5SUMS"
    then
        echo "Unable to download $base_url/MD5SUMS" >&2
        return 1
    fi

    if [ "x$IGNORE_GPG" == "x" ]
    then
        if ! $DOWNLOAD "$base_url/MD5SUMS.gpg"
        then
            echo "Unable to download $base_url/MD5SUMS.gpg" >&2
            return 1
        fi

        if ! gpg --keyring=$GPG_KEYRING --verify MD5SUMS.gpg MD5SUMS >/dev/null 2>/dev/null
        then
            echo "Failed to verify MD5SUMS via $GPG_KEYRING ($base_url/MD5SUMS)" >&2
            return 1
        fi
    fi
    return 0
}

get_md5sum_for_file() {
    local filename=$1

    # The filename supplied in $1 must be the full path as seen in the
    # MD5SUMS file.  The files are rooted from a single place so the grepped
    # string will only match once.
    server_md5sum=$(grep $filename MD5SUMS|awk '{print $1}') ||
        fail "failed to find checksum for $filename"
    echo $server_md5sum
}

check_checksum() {
    local server_md5sum=$1 file_on_disk=$2
    local md5sum

    md5sum=$(md5sum $file_on_disk|awk '{print $1}')

    if [ "$md5sum" != "$server_md5sum" ]; then
        fail "md5 checksum mismatch for $file_on_disk: expected $server_md5sum, got $md5sum"
    fi
}

# Return a list of files for architecture $1 and release $2 that need to be
# downloaded
compose_installer_download_files() {
    local arch=$1 release=$2

    case $arch in
        amd64/*|i386/*)
            echo "linux initrd.gz"
            ;;
        armhf/highbank)
            echo "vmlinuz initrd.gz"
            ;;
        *)
            echo "Unknown architecture: $arch" >&2
            exit 1
            ;;
    esac
}


# Rename downloaded files for architecture $1 and release $2 into the form that
# MAAS expects them
rename_installer_download_files() {
    local arch=$1 release=$2

    case $arch in
        amd64/*|i386/*)
            # do nothing
            ;;
        armhf/highbank)
            mv vmlinuz linux
            ;;
        *)
            echo "Unknown architecture: $arch" >&2
            exit 1
            ;;
    esac
}


# Copy the pre-boot loader pxelinux.0, and modules we need, from the
# installed syslinux version.  Install it into the TFTP tree for
# netbooting.
update_pre_boot_loader() {
    for loader_file in pxelinux.0 chain.c32 ifcpu64.c32
    do
        maas-provision install-pxe-bootloader \
            --loader="/usr/lib/syslinux/$loader_file"
    done
}


# Download kernel/initrd for installing Ubuntu release $2 for
# architecture $1, and install them into the TFTP directory hierarchy.
update_install_files() {
    local arch=$1 release=$2
    local files file url file_prefix filename_in_md5sums_file md5sum

    files=$(compose_installer_download_files $arch $release)
    url=$(compose_installer_download_url $arch $release)

    mkdir "install"
    pushd "install" >/dev/null
    if ! fetch_server_md5sums $(compose_installer_base_url $arch $release)
    then
        echo "Failed to download MD5 sums for $arch $release." >&2
        popd >/dev/null
        rm -rf "install"
        return
    fi
    echo "MD5SUMS GPG signature OK for $arch $release"
    for file in $files
    do
        if ! $DOWNLOAD $url/$file
        then
            # Download failed.  Log error, and skip this image.
            echo "Failed to download $url/$file" >&2
            popd >/dev/null
            rm -rf "install"
            return
        fi
        file_prefix=$(compose_installer_download_url_postfix $arch)
        filename_in_md5sums_file=$file_prefix$file
        md5sum=$(get_md5sum_for_file $filename_in_md5sums_file)
        check_checksum $md5sum $file
        echo "'$file' md5sum OK"
    done
    rename_installer_download_files $arch $release
    popd >/dev/null

    maas-provision install-pxe-image \
        --arch="${arch%%/*}" --subarch="${arch#*/}" \
        --release=$release --purpose="install" \
        --image="install"
}


# Download and install the "install" images.
import_install_images() {
    local arch release DOWNLOAD_DIR

    DOWNLOAD_DIR=$(mktemp -d)
    echo "Downloading to temporary location $DOWNLOAD_DIR."
    pushd -- $DOWNLOAD_DIR

    for arch in $ARCHES
    do
        for release in $RELEASES
        do
            update_install_files $arch $release
        done
    done

    popd
    rm -rf -- $DOWNLOAD_DIR
}


# Download and install the ephemeral images.
import_ephemeral_images() {
    if test "$IMPORT_EPHEMERALS" != "0"
    then
        maas-import-ephemerals
    fi
}


main() {
    # All files we create here are public.  The TFTP user will need to be
    # able to read them.
    umask a+r

    update_pre_boot_loader
    import_install_images
    import_ephemeral_images
}

# check for commandline arguments
if [ $# -gt 0 ]
  then
   case $1 in
    "-h"|"--help") show_usage ; exit ;;
    esac
fi

if [ ! -f "$GPG_KEYRING" ]; then
    fail "gpg keyring $GPG_KEYRING is not a file"
fi

main
