#!/bin/bash

# template script for generating ubuntu container for LXC based on released
# cloud images.
#
# Copyright © 2012 Serge Hallyn <serge.hallyn@canonical.com>
#
#  This library is free software; you can redistribute it and/or
#  modify it under the terms of the GNU Lesser General Public
#  License as published by the Free Software Foundation; either
#  version 2.1 of the License, or (at your option) any later version.

#  This library 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
#  Lesser General Public License for more details.

#  You should have received a copy of the GNU Lesser General Public
#  License along with this library; if not, write to the Free Software
#  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA

set -e

STATE_DIR="/var"
HOOK_DIR="/usr/share/lxc/hooks"
CLONE_HOOK_FN="$HOOK_DIR/ubuntu-cloud-prep"
LXC_TEMPLATE_CONFIG="/usr/share/lxc/config"
KNOWN_RELEASES="precise trusty xenial yakkety zesty"
skip_arch_check=${UCTEMPLATE_SKIP_ARCH_CHECK:-0}

# Make sure the usual locations are in PATH
export PATH=$PATH:/usr/sbin:/usr/bin:/sbin:/bin

if [ -r /etc/default/lxc ]; then
    . /etc/default/lxc
fi

am_in_userns() {
    [ -e /proc/self/uid_map ] || { echo no; return; }
    [ "$(wc -l /proc/self/uid_map | awk '{ print $1 }')" -eq 1 ] || { echo yes; return; }
    line=$(awk '{ print $1 " " $2 " " $3 }' /proc/self/uid_map)
    [ "$line" = "0 0 4294967295" ] && { echo no; return; }
    echo yes
}

in_userns=0
[ $(am_in_userns) = "yes" ] && in_userns=1

copy_configuration()
{
    path=$1
    rootfs=$2
    name=$3
    arch=$4
    release=$5

    if [ $arch = "i386" ]; then
        arch="i686"
    fi

    # if there is exactly one veth network entry, make sure it has an
    # associated hwaddr.
    nics=`grep -e '^lxc\.net\.0\.type[ \t]*=[ \t]*veth' $path/config | wc -l`
    if [ $nics -eq 1 ]; then
        grep -q "^lxc.net.0.hwaddr" $path/config || sed -i -e "/^lxc\.net\.0\.type[ \t]*=[ \t]*veth/a lxc.net.0.hwaddr = 00:16:3e:$(openssl rand -hex 3| sed 's/\(..\)/\1:/g; s/.$//')" $path/config
    fi

    # Generate the configuration file
    ## Relocate all the network config entries
    sed -i -e "/lxc.net.0/{w ${path}/config-network" -e "d}" $path/config

    ## Relocate any other config entries
    sed -i -e "/lxc./{w ${path}/config-auto" -e "d}" $path/config

    ## Add all the includes
    echo "" >> $path/config
    echo "# Common configuration" >> $path/config
    if [ -e "${LXC_TEMPLATE_CONFIG}/ubuntu-cloud.common.conf" ]; then
        echo "lxc.include = ${LXC_TEMPLATE_CONFIG}/ubuntu-cloud.common.conf" >> $path/config
    fi
    if [ -e "${LXC_TEMPLATE_CONFIG}/ubuntu-cloud.${release}.conf" ]; then
        echo "lxc.include = ${LXC_TEMPLATE_CONFIG}/ubuntu-cloud.${release}.conf" >> $path/config
    fi
    if [ $in_userns -eq 1 ] && [ -e "${LXC_TEMPLATE_CONFIG}/ubuntu-cloud.userns.conf" ]; then
        echo "lxc.include = ${LXC_TEMPLATE_CONFIG}/ubuntu-cloud.userns.conf" >> $path/config
    fi

    ## Add the container-specific config
    echo "" >> $path/config
    echo "# Container specific configuration" >> $path/config
    [ -e "$path/config-auto" ] && cat $path/config-auto >> $path/config && rm $path/config-auto
    grep -q "^lxc.rootfs.path" $path/config 2>/dev/null || echo "lxc.rootfs.path = $rootfs" >> $path/config
    cat <<EOF >> $path/config
lxc.uts.name = $name
lxc.arch = $arch
EOF

    ## Re-add the previously removed network config
    echo "" >> $path/config
    echo "# Network configuration" >> $path/config
    cat $path/config-network >> $path/config
    rm $path/config-network

    # Set initial timezone as on host
    if [ -f /etc/timezone ]; then
        cat /etc/timezone > $rootfs/etc/timezone
        chroot $rootfs dpkg-reconfigure -f noninteractive tzdata
    elif [ -f /etc/sysconfig/clock ]; then
        . /etc/sysconfig/clock
        echo $ZONE > $rootfs/etc/timezone
        chroot $rootfs dpkg-reconfigure -f noninteractive tzdata
    else
        echo "Timezone in container is not configured. Adjust it manually."
    fi

    # rmdir /dev/shm for containers that have /run/shm
    # I'm afraid of doing rm -rf $rootfs/dev/shm, in case it did
    # get bind mounted to the host's /run/shm.  So try to rmdir
    # it, and in case that fails move it out of the way.
    # NOTE: This can only be removed once 12.04 goes out of support
    if [ ! -L $rootfs/dev/shm ] && [ -e $rootfs/dev/shm ]; then
        rmdir $rootfs/dev/shm 2>/dev/null || mv $rootfs/dev/shm $rootfs/dev/shm.bak
        ln -s /run/shm $rootfs/dev/shm
    fi

    return 0
}

usage()
{
    cat <<EOF
LXC Container configuration for Ubuntu Cloud images.

Generic Options
[ -r | --release <release> ]: Release name of container, defaults to host
[ --rootfs <path> ]: Path in which rootfs will be placed
[ -a | --arch ]: Architecture of container, defaults to host architecture
[ -T | --tarball ]: Location of tarball
[ -d | --debug ]: Run with 'set -x' to debug errors
[ -s | --stream]: Use specified stream rather than 'tryreleased'

Additionally, clone hooks can be passed through (ie, --userdata).  For those,
see:
  $CLONE_HOOK_FN --help
EOF
    return 0
}

options=$(getopt -o a:hp:r:n:Fi:CLS:T:ds:u: -l arch:,help,rootfs:,path:,release:,name:,flush-cache,hostid:,auth-key:,cloud,no_locales,tarball:,debug,stream:,userdata:,vendordata:,mapped-uid:,mapped-gid: -- "$@")
if [ $? -ne 0 ]; then
    usage $(basename $0)
    exit 1
fi
eval set -- "$options"

mapped_uid=-1
mapped_gid=-1
# default release is trusty, or the systems release if recognized
release=trusty
if [ -f /etc/lsb-release ]; then
    . /etc/lsb-release
    rels=$(ubuntu-distro-info --supported 2>/dev/null) ||
        rels="$KNOWN_RELEASES"
    for r in $rels; do
        [ "$DISTRIB_CODENAME" = "$r" ] && release="$r"
    done
fi

# Code taken from debootstrap
if [ -x /usr/bin/dpkg ] && /usr/bin/dpkg --print-architecture >/dev/null 2>&1; then
    arch=`/usr/bin/dpkg --print-architecture`
elif type udpkg >/dev/null 2>&1 && udpkg --print-architecture >/dev/null 2>&1; then
    arch=`/usr/bin/udpkg --print-architecture`
else
    arch=$(uname -m)
    if [ "$arch" = "i686" ]; then
        arch="i386"
    elif [ "$arch" = "x86_64" ]; then
        arch="amd64"
    elif [ "$arch" = "armv7l" ]; then
        # note: arm images don't exist before oneiric;  are called armhf in
        # trusty and later;  and are not supported by the query, so we don't actually
        # support them yet (see check later on).  When Query2 is available,
        # we'll use that to enable arm images.
        arch="armhf"
    elif [ "$arch" = "aarch64" ]; then
        arch="arm64"
    elif [ "$arch" = "ppc64le" ]; then
        arch="ppc64el"
    fi
fi

debug=0
hostarch=$arch
cloud=0
locales=1
flushcache=0
stream="tryreleased"
cloneargs=()
while true
do
    case "$1" in
    -h|--help)         usage $0 && exit 1;;
    -p|--path)         path=$2; shift 2;;
    -n|--name)         name=$2; shift 2;;
    -F|--flush-cache)  flushcache=1; shift 1;;
    -r|--release)      release=$2; shift 2;;
    -a|--arch)         arch=$2; shift 2;;
    -T|--tarball)      tarball=$2; shift 2;;
    -d|--debug)        debug=1; shift 1;;
    -s|--stream)       stream=$2; shift 2;;
    --rootfs)          rootfs=$2; shift 2;;
    -L|--no?locales)   cloneargs[${#cloneargs[@]}]="--no-locales"; shift 1;;
    -i|--hostid)       cloneargs[${#cloneargs[@]}]="--hostid=$2"; shift 2;;
    -u|--userdata)     cloneargs[${#cloneargs[@]}]="--userdata=$2"; shift 2;;
    -V|--vendordata)   cloneargs[${#cloneargs[@]}]="--vendordata=$2"; shift 2;;
    -C|--cloud)        cloneargs[${#cloneargs[@]}]="--cloud"; shift 1;;
    -S|--auth-key)     cloneargs[${#cloneargs[@]}]="--auth-key=$2"; shift 2;;
    --mapped-uid)      mapped_uid=$2; shift 2;;
    --mapped-gid)      mapped_gid=$2; shift 2;;
    --)                shift 1; break ;;
        *)              break ;;
    esac
done

cloneargs=( "--name=$name" "${cloneargs[@]}" )

if [ $debug -eq 1 ]; then
    set -x
fi

if [ "$arch" = "i686" ]; then
    arch=i386
fi

if [ "$skip_arch_check" = "0" ]; then
    case "$hostarch:$arch" in
        $arch:$arch) : ;; # the host == container
        amd64:i386) :;; # supported "cross"
        arm64:arm*) :;; # supported "cross"
        armel:armhf) :;; # supported "cross"
        armhf:armel) :;; # supported "cross"
        *) echo "cannot create '$arch' container on hostarch '$hostarch'";
           exit 1;;
    esac
fi

if [ "$stream" != "daily" -a "$stream" != "released" -a "$stream" != "tryreleased" ]; then
    echo "Only 'daily' and 'released' and 'tryreleased' streams are supported"
    exit 1
fi

if [ -z "$path" ]; then
    echo "'path' parameter is required"
    exit 1
fi

if [ "$(id -u)" != "0" ]; then
    echo "This script should be run as 'root'"
    exit 1
fi

# detect rootfs
config="$path/config"
if [ -z "$rootfs" ]; then
    if grep -q '^lxc.rootfs.path' $config 2>/dev/null ; then
        rootfs=$(awk -F= '/^lxc.rootfs.path =/{ print $2 }' $config)
    else
        rootfs=$path/rootfs
    fi
fi

type ubuntu-cloudimg-query
type wget

# determine the url, tarball, and directory names
# download if needed
# Allow the cache base to be set by environment variable
cache=${LXC_CACHE_PATH:-"$STATE_DIR/cache/lxc"}/cloud-$release
if [ $in_userns -eq 1 ]; then
    STATE_DIR="$HOME/.cache/lxc"
    cache=${LXC_CACHE_PATH:-"$STATE_DIR"}/cloud-$release
fi

mkdir -p $cache

if [ "$stream" = "tryreleased" ]; then
    stream=released
    ubuntu-cloudimg-query $release $stream $arch 1>/dev/null 2>/dev/null || stream=daily
fi

if [ -n "$tarball" ]; then
    url2="$tarball"
else
    if ! url1=`ubuntu-cloudimg-query $release $stream $arch --format "%{url}\n"`; then
        echo "There is no download available for release=$release, stream=$stream, arch=$arch"
        [ "$stream" = "daily" ] || echo "You may try with '--stream=daily'"
        exit 1
    fi
    if [ "$release" = "precise" ] || [ "$release" = "trusty" ]; then
        url2=`echo $url1 | sed -e 's/.tar.gz/-root\0/' -e 's/.tar.gz/.tar.xz/'`
    else
        url2=`echo $url1 | sed -e 's/.tar.gz/.squashfs/'`
    fi
fi

filename=`basename $url2`

wgetcleanup()
{
    rm -f $filename
}

do_extract_rootfs() {

    cd $cache
    if [ $flushcache -eq 1 ]; then
        echo "Clearing the cached images"
        rm -f $filename
    fi

    trap wgetcleanup EXIT SIGHUP SIGINT SIGTERM
    if [ ! -f $filename ]; then
        wget $url2
    fi
    trap EXIT
    trap SIGHUP
    trap SIGINT
    trap SIGTERM

    echo "Extracting container rootfs"
    mkdir -p $rootfs
    cd $rootfs
    if [ "${filename##*.}" = "squashfs" ]; then
        unsquashfs -n -f -d "$rootfs" "$cache/$filename"
    else
        if [ $in_userns -eq 1 ]; then
            tar --anchored --exclude="dev/*" --numeric-owner -xpf "$cache/$filename"
            mkdir -p $rootfs/dev/pts/
        else
            tar --numeric-owner -xpf "$cache/$filename"
        fi
    fi
}

if [ -n "$tarball" ]; then
    do_extract_rootfs
else
    mkdir -p "$STATE_DIR/lock/subsys/"
    (
        flock -x 9
        do_extract_rootfs
    ) 9>"$STATE_DIR/lock/subsys/lxc-ubuntu-cloud"
fi

copy_configuration $path $rootfs $name $arch $release

"$CLONE_HOOK_FN" "${cloneargs[@]}" "$rootfs"

if [ $mapped_uid -ne -1 ]; then
    chown $mapped_uid $path/config
    chown -R $mapped_uid $STATE_DIR
    chown -R $mapped_uid $cache
fi
if [ $mapped_gid -ne -1 ]; then
    chgrp $mapped_gid $path/config
    chgrp -R $mapped_gid $STATE_DIR
    chgrp -R $mapped_gid $cache
fi

echo "Container $name created."
exit 0

# vi: ts=4 expandtab
