#!/bin/bash

# (C) Copyright Canonical 2011,2012

# What lxc container to clone
LXC_BASE=""
# $2 is a path to bind mount e.g. /tmp/foo.
LXC_BIND=""
UNION="overlayfs"

usage() {
    echo "usage: lxc-start-ephemeral [-U uniontype] [-d|--daemon] [-h] [-b bdir] [-u user] [-S key] -o orig -- [COMMAND [ARGS...]]"
}

help() {
    usage
    echo
    echo "Runs an ephemeral (one-off) container"
    echo
    echo "Options:"
    echo "orig        : name of the original container"
    echo "bdir        : directory to bind mount into container"
    echo "user        : the user to connect to the container as"
    echo "key         : the path to the SSH key to use to connect"
    echo "-d          : run in the background"
    echo "-U          : type of union (aufs or overlayfs)"
    echo "              Default is overlayfs"
    echo
    echo "if a COMMAND is given, then the container will run only as long"
    echo "as the command runs.  If no COMMAND is given, this command will"
    echo "wait until the container is shut down"
}

shortoptions='hb:o:u:S:dU:'
longoptions='help,orig:,bdir:,user:,ssh-key:,daemon,union:'

LXC_RUNNING=0
LXC_MOUNTED=0
DAEMON=0

cleanup() {
    if [ $LXC_RUNNING -eq 1 ]; then
        sudo lxc-stop -n $LXC_NAME
    fi
    if [ $LXC_MOUNTED -eq 1 ]; then
        sudo umount $EPHEMERAL_BIND_DIR
        sudo umount $LXC_DIR
        sudo umount $OVERLAY_DIR
        sudo rmdir $LXC_DIR
        sudo rmdir $OVERLAY_DIR
    fi
    exit 1
}

do_mount() {
    lower=$1
    upper=$2
    target=$3
    if [ $UNION = "aufs" ]; then
        sudo mount -t aufs -o br=${upper}=rw:${lower}=ro,noplink none ${target}
    else
        sudo mount -t overlayfs -oupperdir=${upper},lowerdir=${lower} none ${target}
    fi
}

trap cleanup SIGTERM SIGINT SIGQUIT 

getopt=$(getopt -o $shortoptions --longoptions  $longoptions -- "$@")
if [ $? != 0 ]; then
    usage
    exit 1;
fi

eval set -- "$getopt"

while true; do
        case "$1" in
	    -h|--help)
		help
		exit 1
		;;
	    -o|--orig)
		shift
		LXC_BASE=$1
		shift
		;;
	    -b|--bdir)
		shift
		LXC_BIND=$1
		shift
		;;
	    -u|--user)
		shift
		LXC_USER=$1
		shift
		;;
	    -S|--ssh-key)
		shift
		LXC_KEY="-i $1"
		shift
		;;
            -d|--detach)
		DAEMON=1
		shift
		;;
	    -U|--union)
		shift
		UNION=$1
		shift
		;;
            --)
		shift
		break;;
            *)
		echo $1
		usage
		exit 1
		;;
        esac
done

COMMAND=$@
COMMAND_LENGTH=$#
LXC_USER=${LXC_USER:-`id -un`}

# validation

if [ -z $LXC_BASE ]; then
    echo "original container must be specified"
    usage
    exit 1
fi
if [ ! -d /var/lib/lxc/$LXC_BASE ] ; then
    echo 'no such lxc container $LXC_BASE'
    exit 1
fi

if [ "$UNION" != "overlayfs" -a "$UNION" != "aufs" ]; then
    echo "Invalid option for union: choose overlayfs or aufs."
    exit 1
fi

setup_container()
{
    echo "Setting up ephemeral container..."
    OVERLAY_DIR=`mktemp -d /tmp/lxc-lp-XXXXXXX`
    sudo mount -t tmpfs none $OVERLAY_DIR
    LXC_DIR=`sudo mktemp -d --tmpdir=/var/lib/lxc $LXC_BASE-temp-XXXXXXX`
    LXC_NAME=`basename $LXC_DIR`
    do_mount "/var/lib/lxc/$LXC_BASE" "${OVERLAY_DIR}" $LXC_DIR
    EPHEMERAL_BIND_DIR=$LXC_DIR/ephemeralbind
    sudo mkdir $EPHEMERAL_BIND_DIR
    sudo mount -t tmpfs none $EPHEMERAL_BIND_DIR
    LXC_MOUNTED=1

    # Update the ephemeral lxc's configuration to reflect the new container name.
    sudo sed -i -e "s/$LXC_BASE/$LXC_NAME/" $LXC_DIR/fstab $LXC_DIR/config $LXC_DIR/rootfs/etc/hostname $LXC_DIR/rootfs/etc/hosts

    # Update the fstab to have all bind mounts be ephemeral.
    sudo cp $LXC_DIR/fstab $LXC_DIR/fstab.old
    while read line; do
        # Pull out the second field of the current line of fstab info
        path=`echo -n $line | awk '{print $2}'`
        # If LXC_BIND is not set, or the mount destination of this line is not
        # LXC_BIND...
        if [ -z "$LXC_BIND" -o "`readlink -f $path`" != "`readlink -f $LXC_DIR/rootfs$LXC_BIND`" ];
            then
            # ...then we should write some form of this line.
            # If this line is a bind...
            if [ `echo -n $line | awk '{print $4}'` == "bind" ]; then
                # ...we should rewrite it as an overlay.
                source=`echo -n $line | awk '{print $1}'`
                upperdir=$EPHEMERAL_BIND_DIR$source
                sudo mkdir -p $upperdir
                sudo chown `sudo stat -c '%U.%G' $source` $upperdir
                if [ $UNION = "overlayfs" ]; then
                    echo "none	$path	overlayfs upperdir=$upperdir,lowerdir=$source 0 0";
                else
                    echo "none $path aufs br=${upperdir}=rw:${lowerdir}=ro,noplink 0 0";
                fi
            else
                # Otherwise, we can pass it through unchanged.
                echo "$line";
            fi
        fi
    done < $LXC_DIR/fstab.old > $LXC_DIR/fstab

    # If LXC_BIND is defined, add it to fstab.
    if [ -n "$LXC_BIND" ]; then
        sudo mkdir -p $LXC_DIR/rootfs$LXC_BIND
        echo "$LXC_BIND	$LXC_DIR/rootfs$LXC_BIND	none bind 0 0" >> $LXC_DIR/fstab
    fi

    # update the ephemeral container's MAC address (lifted from lxc-clone)
    c=$LXC_DIR/config
    # change hwaddrs
    mv ${c} ${c}.old
    (
    while read line; do
        if [ "${line:0:18}" = "lxc.network.hwaddr" ]; then
                echo "lxc.network.hwaddr= 00:16:3e:$(openssl rand -hex 3| sed 's/\(..\)/\1:/g; s/.$//')"
        else
                echo "$line"
        fi
    done
    ) < ${c}.old > ${c}
    rm -f ${c}.old

    # precise is now the worst - its dhclient *uses*
    # /var/lib/dhcp3/dhclient.eth0.leases but ships without that
    # directory, so you get no leases file at all.
    LEASES=$LXC_DIR/rootfs/var/lib/dhcp3/dhclient.eth0.leases
    LEASES2=$LXC_DIR/rootfs/var/lib/dhcp/dhclient.leases
    sudo mkdir -p $LXC_DIR/rootfs/var/lib/dhcp3
    sudo mkdir -p $LXC_DIR/rootfs/var/lib/dhcp
    sudo truncate -c -s0 $LEASES
    sudo truncate -c -s0 $LEASES2
}

start_container()
{
    echo "Starting up the container..."
    sudo lxc-start -n $LXC_NAME -d
    LXC_RUNNING=1

    if [ $COMMAND_LENGTH -gt 0 ]; then
        # when lxc-attach support arrives in the kernel, we can switch to
        # that
        delay=30
        while [ $delay -gt 0 -a ! -s $LEASES -a ! -s $LEASES2 ]
        do
            delay=$(( $delay - 1 ))
            sleep 1
        done

        [ -s $LEASES ] || LEASES=$LEASES2
        IP_ADDRESS=`sudo grep fixed-address $LEASES | tail -n 1 | sed -r 's/.* ([0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}).*/\1/'`
        if [ 0 -eq $? -a -n "$IP_ADDRESS" ]; then
            # Repeatedly try to connect over SSH until we either succeed
            # or time out.
            for i in $(seq 1 30); do
            ssh -n -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null $LXC_KEY $LXC_USER@$IP_ADDRESS -- "$COMMAND"
                if [ ! 255 -eq $? ]; then
                    # If ssh returns 255 then its connection failed.
                    # Anything else is either success (status 0) or a
                    # failure from whatever we ran over the SSH connection.
                    # In those cases we want to stop looping, so we break
                    # here.
                    break;
                fi
                sleep 1
            done
        else
            echo "could not get IP address - aborting." >&2
            echo "content of $LEASES:" >&2
            cat $LEASES >&2
            echo "content of $LEASES2:" >&2
            cat $LEASES2 >&2
        fi
    else
        sudo lxc-wait -n $LXC_NAME -s RUNNING
        echo "$LXC_NAME is running"
        echo "You connect with the command:"
        echo "    sudo lxc-console -n $LXC_NAME"
        sudo lxc-wait -n $LXC_NAME -s STOPPED
    fi
}

stop_container()
{
    echo "Stopping lxc" >&2
    sudo lxc-stop -n $LXC_NAME
    sleep 2
    # echo "umounting ephemeral_bind_dir $EPHEMERAL_BIND_DIR" >&2
    sudo umount $EPHEMERAL_BIND_DIR
    # echo "umounting lxc_dir $LXC_DIR" >&2
    sudo umount $LXC_DIR
    # echo "umounting overlay" >&2
    sudo umount $OVERLAY_DIR
    # echo "rming lxc_dir $LXC_DIR" >&2
    sudo rmdir $LXC_DIR
    # echo "rming overlay dir $OVERLAY_DIR" >&2
    sudo rmdir $OVERLAY_DIR
}

handle_container()
{
    setup_container
    start_container
    stop_container
    exit 0
}

if [ $DAEMON -eq 1 ]; then
    handle_container &
    exit 0
fi

handle_container
