#!/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_dirs()
{
    # 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/rootfs
    # echo "umounting overlay" >&2
    sudo umount $OVERLAY_DIR

    # remove all contents of the content dir
    sudo rmdir $LXC_DIR/ephemeralbind
    sudo rmdir $LXC_DIR/rootfs
    sudo rm -f $LXC_DIR/*

    # echo "rming lxc_dir $LXC_DIR" >&2
    sudo rmdir $LXC_DIR
    # echo "rming overlay dir $OVERLAY_DIR" >&2
    sudo rmdir $OVERLAY_DIR
}

cleanup() {
    if [ $LXC_RUNNING -eq 1 ]; then
        sudo lxc-stop -n $LXC_NAME
    fi
    if [ $LXC_MOUNTED -eq 1 ]; then
        cleanup_dirs
    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`
    sudo chmod 755 ${LXC_DIR}
    LXC_NAME=`basename $LXC_DIR`
    sudo mkdir ${LXC_DIR}/rootfs
    do_mount "/var/lib/lxc/$LXC_BASE/rootfs" "${OVERLAY_DIR}" ${LXC_DIR}/rootfs
    EPHEMERAL_BIND_DIR=$LXC_DIR/ephemeralbind
    sudo mkdir $EPHEMERAL_BIND_DIR
    sudo mount -t tmpfs none $EPHEMERAL_BIND_DIR
    LXC_MOUNTED=1
    {
        d1=/var/lib/lxc/${LXC_BASE}
        for f in ${d1}/*; do
            if [ -f $f ]; then
                sudo cp $f $LXC_DIR/
            fi
        done
    }

    # Update the ephemeral lxc's configuration to reflect the new container name.
    # Check all the places known distros keep hostnames.
    # FIXME: should we sanity check the hostname to make sure it contains no bad chars?
    for file in $LXC_DIR/fstab $LXC_DIR/config \
                $LXC_DIR/rootfs/etc/hostname \
                $LXC_DIR/rootfs/etc/hosts \
                $LXC_DIR/rootfs/etc/sysconfig/network \
                $LXC_DIR/rootfs/etc/sysconfig/network-scripts/ifcfg-eth0
    do
        if test -f "$file"
        then
            sudo sed -i -e "s/$LXC_BASE/$LXC_NAME/" $file
        fi
    done

    # 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 | sudo tee $LXC_DIR/fstab >/dev/null

    # 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" | sudo tee -a $LXC_DIR/fstab >/dev/null
    fi

    # update the ephemeral container's MAC address (lifted from lxc-clone)
    c=$LXC_DIR/config
    # change hwaddrs
    sudo 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 | sudo tee ${c} >/dev/null
    sudo rm -f ${c}.old
}

get_ip()
{
    # Get init's PID
    PID=$(sudo lxc-info -n $1 -p | awk '{print $2}')
    [ "$PID" = "-1" ] && return 1

    # Get some unique path
    DST=$(sudo mktemp -u --tmpdir=/run/netns/)
    NAME=$(basename $DST)

    # Prepare the /run/netns entry for "ip netns"
    sudo mkdir -p /run/netns
    sudo ln -s /proc/$PID/ns/net $DST

    # Grab all the public globally routed IPv4 and IPv6 addresses
    (sudo ip netns exec $NAME ip -4 addr show scope global && \
     sudo ip netns exec $NAME ip -6 addr show scope global) | grep inet | while read line; do
        ip=$(echo $line | awk '{print $2}' | cut -d '/' -f1)
        echo "$ip"
    done

    sudo rm $DST
}

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

    if [ $COMMAND_LENGTH -gt 0 ]; then
        # When lxc-attach support arrives in the kernel, we can switch to
        # that.
        # Meanwhile, we use get_ip to wait for container's network to be up
        # and to obtain the ip address, then we can ssh to the lxc.
        TRIES=60
        FAILED=1

        # Repeatedly try to connect over SSH until we either succeed
        # or time out.
        for i in $(seq 1 $TRIES); do
            # We call get_ip inside the loop to ensure the correct ip
            # is retrieved even in the case the DHCP ip assignment
            # changes during the process.
            IP_ADDRESS=$(get_ip $LXC_NAME)
            if [ -z "$IP_ADDRESS" ]; then
                sleep 1
                continue
            fi

            # Iterate through all the addresses (if multiple)
            for ip in $IP_ADDRESS; 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

                    FAILED=0
                    break;
                fi
            done

            if [ "$FAILED" = "0" ]; then
                break
            fi
            sleep 1
        done

        if [ "$FAILED" = "1" ]; then
            echo "could not get IP address - aborting." >&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
    LXC_RUNNING=0
    cleanup_dirs
}

handle_container()
{
    setup_container
    start_container
    stop_container
    exit 0
}

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

handle_container
