#!/bin/bash
#
# maas-import-ephemerals - sync and import ephemeral images
#
# Copyright (C) 2011-2012 Canonical
#
# Authors:
#    Scott Moser <scott.moser@canonical.com>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, version 3 of the License.
#
# 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 Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

VERBOSITY=0

# Mirror to load cloud images from.  When the cluster controller runs the
# import scripts, it provides a setting from the server side.
CLOUD_IMAGES_ARCHIVE=${CLOUD_IMAGES_ARCHIVE:-https://maas.ubuntu.com/images}

# iSCSI targets configuration file.
SYS_TGT_CONF="/etc/tgt/targets.conf"

# Prefix for iSCSI target name.
TARGET_NAME_PREFIX="iqn.2004-05.com.ubuntu:maas:"

# TODO: What's this for?  If set, it gets run on a downloaded disk.img,
# kernel, and initrd.
EPH_UPDATE_CMD=""

# TODO: DATA_DIR seems to be the root of a directory tree that's exposed over
# iSCSI for download by nodes.  Can we confirm this?
DATA_DIR="/var/lib/maas/ephemeral"

# Optional configuration script that may set variables for use by this
# script.  It gets sourced later on.
CONFIG="/etc/maas/import_ephemerals"

RELEASES="precise"
ARCHES="amd64/generic i386/generic armhf/highbank"
BUILD_NAME="ephemeral"
STREAM="released"

# DATA_DIR layout is like:
#   tgt.conf
#   tgt.conf.d/
#     <name>.conf ->
#        ../release/stream/arch/serial.conf
#   release/
#     stream/
#       arch/
#         serial/
#           kernel
#           disk.img
#           initrd
#           my.conf

error() { echo "$@" 1>&2; }
fail() { [ $# -eq 0 ] || error "$@"; exit 1; }


Usage() {
    cat <<EOF
Usage: ${0##*/} [ options ] <<ARGUMENTS>>

   Import ephemeral (enlistment/commissioning) images into maas.
   Settings are read from $CONFIG.
EOF
}


bad_Usage() { Usage 1>&2; [ $# -eq 0 ] || error "$@"; exit 1; }


cleanup() {
    [ -z "${TEMP_D}" -o ! -d "${TEMP_D}" ] || rm -Rf "${TEMP_D}"
}


debug() {
    local level=${1}; shift;
    [ "${level}" -gt "${VERBOSITY}" ] && return
    error "${@}"
}


uniq_major_arches() {
    # print on stdout a uniq set of arches out of the list of arch/subarch
    # pairs supplied in $@
    # Eg. armhf/highbank armhf/armadaxp i386/generic amd64/generic ->
    #           "armhf\ni386\namd64\n"
    local arch
    for arch in "$@"; do echo "${arch%%/*}"; done|uniq
}


subarches() {
    # print on stdout a list of subarches available for a given major arch $1
    # given a list of arch/subarch pairs in $2-
    # Eg. armhf armhf/highbank armhf/armadaxp i386/generic ->
    #           "highbank\narmadaxp\n"
    local major_arch="$1" candidate
    shift
    for candidate in "$@"; do
        case "$candidate" in
            "$major_arch"/*) echo "${candidate#*/}" ;;
            *) ;;
        esac
    done
}


query_remote() {
    # query /query data at CLOUD_IMAGES_ARCHIVE
    # returns 7 values prefixed with 'r_'
    local iarch=$1 irelease=$2 istream=$3 out=""
    local burl="${CLOUD_IMAGES_ARCHIVE}/query"
    local url="$burl/$irelease/$istream/${STREAM}-dl.current.txt"
    local target="$TEMP_D/query/$release.$stream"
    mkdir -p -- "$TEMP_D/query"
    if [ ! -f "$TEMP_D/query/$release.$stream" ]; then
        wget -q "$url" -O "$target.tmp" && mv "$target.tmp" "$target" ||
            { error "failed to get $url"; return 1; }
    fi

    r_release=""; r_stream=""; r_label=""; r_serial="";
    r_arch=""; r_url=""; r_name=""

    out=$(awk '-F\t' '$1 == release && $2 == stream && $5 == arch { print $3, $4, $6, $7 }' \
        "arch=$iarch" "release=$irelease" "stream=$istream" \
        "$target") && [ -n "$out" ] ||
        return 1

    set -- ${out}
    r_release=$irelease
    r_stream=$istream
    r_label=$1;
    r_serial=$2;
    r_arch=$iarch
    r_url=$3
    r_name=$4
    return
}


query_local() {
    local iarch=$1 irelease=$2 istream=$3 out=""
    local label="" name="" serial="" url=""

    local found=""
    for i in "${DATA_DIR}/"$irelease/$istream/$iarch/*/info; do
        [ -f "$i" ] && found=$i
    done
    found=$(LC_ALL=C;
        cd "${DATA_DIR}/$irelease/$istream/$iarch" 2>/dev/null || exit 0;
        for d in [0-9][0-9][0-9][0-9][0-9][0-9][0-9][0-9]*; do
            [ -f "$d/info" ] && f=$d; done;
        [ -n "$f" ] && echo "$PWD/$f/info")

    l_release=""; l_stream=""; l_label=""; l_serial="";
    l_arch=""; l_url=""; l_name=""; l_dir="";
    if [ -n "$found" ]; then
        . "$found"
        l_release="$release";
        l_stream="$stream";
        l_label="$label";
        l_serial="$serial";
        l_arch="$arch";
        l_url="$url";
        l_name="$name";
        l_dir="${found%/*}";
    fi
}


serial_gt() {
    # is $1 a larger serial than $2 ?
    local a=${1:-0} b=${2:-0}
    case "$a" in
        *.[0-9]) a="${a%.*}${a##*.}";;
    esac
    case "$b" in
        *.[0-9]) b="${b%.*}${b##*.}";;
    esac
    [ $a -gt $b ]
}


prep_dir() {
    local wd="$1" exdir="" tarball=""
    shift
    local release=$1 stream=$2 label=$3 serial=$4 arch=$5 url=$6 name=$7
    local furl="$CLOUD_IMAGES_ARCHIVE/$url"

    mkdir -p "$wd"
    cat > "$wd/info" <<EOF
release=$release
stream=$stream
label=$label
serial=$serial
arch=$arch
url=$url
name=$name
EOF

    # download
    local cachepath="${TARBALL_CACHE_D}/${name}.tar.gz" rmtar=""
    if [ -f "$cachepath" ]; then
        tarball="${cachepath}"
    elif [ -n "$TARBALL_CACHE_D" ]; then
        mkdir -p "$TARBALL_CACHE_D"
        debug 1 "downloading $name from $furl to local cache"
        wget "$furl" --progress=dot:mega -O "${cachepath}.part$$" &&
            mv "$cachepath.part$$" "$cachepath" || {
            rm "$cachepath.part$$"
            error "failed to download $furl";
            return 1;
        }
        tarball="${cachepath}"
    else
        debug 1 "downloading $name from $furl"
        tarball="$wd/dist.tar.gz"
        wget "$furl" --progress=dot:mega -O "${tarball}" ||
            { error "failed to download $furl"; return 1; }
        rmtar="$tarball"
    fi

    # Extract the tarball.
    exdir="$wd/.xx"
    mkdir -p "$exdir" &&
        debug 1 "extracting tarball" &&
        tar -Sxzf - -C  "$exdir" < "$tarball" ||
        { error "failed to extract tarball from $furl"; return 1; }

    # Look for our files in the extracted tarball.
    local x="" img="" kernel="" initrd=""
    for x in "$exdir/"*.img; do
        [ -f "$x" ] && img="$x" && break
    done

    for x in "$exdir/kernel" "$exdir/"*-vmlinuz*; do
        [ -f "$x" ] && kernel="$x" && break
    done

    for x in "$exdir/initrd" "$exdir/"*-initrd*; do
        [ -f "$x" ] && initrd="$x" && break
    done

    # Rename/move files extracted from tarballs to the target dir.
    [ -n "$img" ] || { error "failed to find image in $furl"; return 1; }
    mv "$img" "$wd/disk.img" ||
        { error "failed to move extracted image to $wd/disk.img"; return 1; }

    [ -z "$kernel" ] || mv "$kernel" "$wd/linux" ||
        { error "failed to move extracted kernel to $wd/linux"; return 1; }

    [ -z "$initrd" ] || mv "$initrd" "$wd/initrd.gz" ||
        { error "failed to move extracted initrd to $wd/initrd.gz"; return 1; }

    [ ! -d "$exdir/subarch" ] || mv "$exdir/subarch" "$wd/" ||
        { error "failed to move extracted subarch to $wd/subarch"; return 1; }

    rm -Rf "$exdir" || { error "failed to cleanup extraction dir"; return 1; }
    { [ -z "$rmtar" ] || rm "$rmtar"; } ||
        { error "failed to remove temporary tarball $rmtar"; return 1; }

    if [ -n "$EPH_UPDATE_CMD" ]; then
        # update
        debug 1 "invoking: ${EPH_UPDATE_CMD[*]} ./disk.img ./kernel ./initrd"
        "${EPH_UPDATE_CMD[@]}" "$wd/disk.img" "$wd/kernel" "$wd/initrd" ||
            { error "failed to apply updates to $img"; return 1; }
    else
        [ -n "$kernel" -a -n "$initrd" ] || {
            error "missing kernel or initrd in tarball. set \$EPH_UPDATE_CMD";
            # TODO: Set it to what!?
            return 1;
        }
    fi

    return 0
}


write_tgt_conf() {
    local file="$1" target_name="$2" image="$3"
    shift 2;
    local release=$1 stream=$2 label=$3 serial=$4 arch=$5 url=$6 name=$7
    cat > "$file" <<EOF
<target ${target_name}>
    readonly 1
    backing-store "$image"
</target>
EOF
}


copy_first_available() {
    # Copy file $1 or $2 (the first that is available) to destination $3.
    local preferred_file="$1" alternate_file="$2" destination="$3"
    local actual=""

    if [ -f "${preferred_file}" ]; then
        actual="${preferred_file}"
    elif [ -f "${alternate_file}" ]; then
        actual="${alternate_file}"
    else
        error "Could not copy to ${destination}."
        error "Neither ${preferred_file} nor ${alternate_file} exists."
        return 1
    fi

    cp -- "${actual}" "${destination}" ||
        { error "Could not copy ${actual} to ${destination}."; return 1; }
    return 0
}


install_tftp_image() {
    # Make image in directory $1, for architecture $2 and subarchitecture $3,
    # and OS release $4, available over TFTP for netbooting nodes.  Only the
    # kernel and initrd are needed.

    local src="$1" arch="$2" subarch="$3" release="$4" tmpdir=""

    # Create image in a temporary directory; the installation process
    # deletes it.
    tmpdir="$(mktemp -d)"

    if [ -f "$src/subarch/$subarch/linux" -a \
            -f "$src/subarch/$subarch/initrd.gz" ]; then
        cp "$src/subarch/$subarch/linux" "$tmpdir/linux" || return 1
        cp "$src/subarch/$subarch/initrd.gz" "$tmpdir/initrd.gz" || return 1
    else
        copy_first_available "$src/linux" "$src/kernel" "$tmpdir/linux" ||
            return 1
        copy_first_available "$src/initrd.gz" "$src/initrd" "$tmpdir/initrd.gz" ||
            return 1
    fi

    local cmd out=""
    cmd=( maas-provision install-pxe-image
          "--arch=$arch" "--subarch=$subarch" "--release=$release"
          --purpose="commissioning" --image="$tmpdir" )
    debug 2 "${cmd[@]}"
    out=$("${cmd[@]}" 2>&1) ||
        { error "cmd failed:" "${cmd[@]}"; error "$out"; return 1; }
}


short_opts="hciuv"
long_opts="help,verbose"
getopt_out=$(getopt --name "${0##*/}" \
    --options "${short_opts}" --long "${long_opts}" -- "$@") &&
    eval set -- "${getopt_out}" ||
    bad_Usage

while [ $# -ne 0 ]; do
    cur=${1}; next=${2};
    case "$cur" in
        -h|--help) Usage ; exit 0;;
        -v|--verbose) VERBOSITY=$((${VERBOSITY}+1));;
        --) shift; break;;
    esac
    shift;
done

[ ! -f "$CONFIG" ] || . "$CONFIG"
[ ! -f ".${CONFIG}" ] || . ".${CONFIG}"


mkdir -p "$DATA_DIR" "$DATA_DIR/.working" ||
    fail "failed to make $DATA_DIR"

TEMP_D=$(mktemp -d "$DATA_DIR/.working/${0##*/}.XXXXXX") ||
   fail "failed to make tempdir"
trap cleanup EXIT

tgt_conf_d="$DATA_DIR/tgt.conf.d"
tgt_conf="${DATA_DIR}/tgt.conf"

mkdir -p "$tgt_conf_d" ||
    fail "failed to make directories"
if [ ! -f "${tgt_conf}" ]; then
    cat > "${tgt_conf}" <<EOF
include ${DATA_DIR}/tgt.conf.d/*.conf
default-driver iscsi
EOF
fi

updates=0
for release in $RELEASES; do
    for arch in $(uniq_major_arches $ARCHES); do
        query_local "$arch" "$release" "$BUILD_NAME" ||
            fail "failed to query local for $release/$arch"
        query_remote "$arch" "$release" "$BUILD_NAME" ||
            fail "remote query of $CLOUD_IMAGES_ARCHIVE failed"

        info="rel: $r_release, arch: $arch: name: $r_name"
        debug 2 "$info"
        debug 2 "local  serial=$l_serial l_name=$l_name dir=$l_dir"
        debug 2 "remote serial=$r_serial r_name=$r_name url=$r_url"

        # if remote is newer, need to update.  Note that if there is no local
        # data, 'l_serial' will be "", which serial_gt considers zero
        if serial_gt "$r_serial" "$l_serial"; then
            # an update is needed remote serial is newer than local
            updates=$(($updates+1))

            msg="updating [${l_name:+${l_name} to }$r_name]"
            debug 0 "$release/$arch: $msg"
            wd="${TEMP_D}/$release/$arch"
            prep_dir "$wd" \
                "$r_release" "$r_stream" "$r_label" \
                "$r_serial" "$r_arch" "$r_url" "$r_name" ||
                fail "failed to prepare image for $release/$arch"

            final_d="${r_release}/${r_stream}/${r_arch}/${r_serial}"
            fpfinal_d="${DATA_DIR}/${final_d}"
            mkdir -p "${fpfinal_d}"

            mv "$wd/"* "${fpfinal_d}/" ||
                fail "failed to move contents to final directory ${fpfinal_d}"
            name="${r_name}"
        else
            debug 0 "$release/$arch: up to date [$l_name]"

            fpfinal_d="${l_dir}"
            final_d="${l_release}/${l_stream}/${l_arch}/${l_serial}"

            name="${l_name}"
        fi

        for subarch in $(subarches "$arch" $ARCHES); do
                # Even if there was no need to update the image, we make sure
                # it gets installed.
                debug 1 "adding images for $release/$arch/$subarch to maas"
                install_tftp_image "$fpfinal_d" "$arch" "$subarch" "$release" ||
                    fail "failed to install tftp image [$info]"
        done

        target_name="${TARGET_NAME_PREFIX}${name}"
        rel_tgt="../${final_d}/tgt.conf"

        # iscsi_update
        write_tgt_conf "${fpfinal_d}/tgt.conf" "$target_name" \
            "${fpfinal_d}/disk.img" ||
            fail "failed to write tgt.conf for $release/$arch"

        ln -sf "$rel_tgt" "${tgt_conf_d}/${name}.conf" ||
            fail "failed to symlink ${name}.conf into place"

        ver_out="${TEMP_D}/verify.${target_name}"
        tgt-admin --conf "$SYS_TGT_CONF" --update "${target_name}" &&
            tgt-admin --conf "$SYS_TGT_CONF" --show > "${ver_out}" &&
            grep -q "^Target [0-9][0-9]*: ${target_name}" "${ver_out}" || {
            mv "${fpfinal_d}/info" "${fpfinal_d}/info.failed"
            tgt-admin --conf "$SYS_TGT_CONF" --delete "$target_name"
            rm "${tgt_conf_d}/${name}.conf"
            fail "failed tgt-admin add for $name"
        }

    done
done


## cleanup
# here, go through anything non-current,
#   * remove the tgt config
#   * if tgt-show has entry:
#     * remove from tgt-admin by name && remove directories
#   * else
#     * remove directory

# vi: ts=4 expandtab
