#!/bin/bash

set -e
set -o pipefail

me=$(basename $0)

# Main entry point, called at the end of the script
main()
{
    # We either need to be root or have password-less sudo access to run as root.  Make sure of it.
    if [ $EUID -ne 0 ] ; then
        if ! sudo -n true >/dev/null 2>&1 ; then
	        echo "$me: error: This script must be run as root or as a user with nopassword sudo authority" >&2
	        return 1
        fi
        echo "$me: Running as $USER, re-running with sudo"
        exec sudo -E bash "$0" "$@"
        return 0
    else
        echo "$me: Runnig as root"
    fi

    # Parse command line
    COMMAND=()
    for arg in "$@" ; do
        case "$next" in
            workspace) WORKSPACE="$arg" ; next= ; continue ;;
            preserve) PRESERVE="$arg" ; next= ; continue ;;
            reffile) REF_FILE="$arg" ; next= ; continue ;;
        esac
        case "$arg" in
            -p) next=preserve ;;
            -r) next=reffile ;;
            -w) next=workspace ;;
            *) COMMAND+=("$arg") ;;
        esac
    done

    # Make sure we have any required variable
    if [ -z "$COMMAND" ] ; then
        echo "$me: error: No commands passed on the command line" >&2
        usage
        return 1
    elif [ -z "$WORKSPACE" ] ; then
        echo "$me: error: The WORKSPACE environment variable is not defined or specified with -w" >&2
        usage
        return 1
    elif ! command -v rsync >/dev/null 2>&1 ; then
        echo "$me: error: rsync is not installed" >&2
        return 1
    fi
    if [ -n "$REF_FILE" ] ; then
        REF_UID=$(stat -c %u "$REF_FILE")
        REF_GID=$(stat -c %g "$REF_FILE")
    fi

    for v in WORKSPACE PRESERVE REF_FILE REF_UID REF_GID ; do
        echo "$me: $v: ${!v}"
    done

    # Given WORKSPACE, create a new directory there to play in
    for ((i=0; i < 100; ++i)) ; do
        if out=$(mkdir "$WORKSPACE/chroot_$i" 2>&1) ; then
            BASE="$WORKSPACE/chroot_$i"
            add_exit_trap "rm -rf $BASE"
            break
        fi
    done
    if [ -z "$BASE" ] ; then
        echo "$me: error: Unable to create chroot directory in $WORKSPACE: $out" >&2
        return 1
    fi

    # Set up a directory to mount the overlayfs on
    ROOT=$BASE/root

    # We're going to run things from the current directory, but in the chroot
    RUNDIR=$(pwd)

    mkdir $ROOT

    # Mount / (ro) with an overlay (rw) to our new root
    mount_overlay / $ROOT $BASE

    # Handle a few well-known filesystem mounts point that might be needed
    for mp in /boot /var /usr ; do
	    if mountpoint -q $mp ; then
	        mount_overlay $mp $ROOT$mp $BASE
	    fi
    done

    # Walk up the path from / to RUNDIR to make sure none of those are mount point
    local dirs=()
    local d="$RUNDIR"
    while [ "$d" != "/" ] ; do
        dirs+=("$d")
        d="$(dirname "$d")"
    done
    for ((i=${#dirs[@]}; i >= 0; --i)) ; do
        d="${dirs[i]}"
        if mountpoint -q "$d" ; then
            mount_overlay "$d" $ROOT$d $BASE
        fi
    done

    # Bind mount /dev and /proc into the new root since we can't overlay these types of filesystems
    # It would be nice to then remount these read-only, but things like /dev/null need to be writable
    mount --bind /proc $ROOT/proc
    add_exit_trap "umount $ROOT/proc"
    mount --bind /dev $ROOT/dev
    add_exit_trap "umount $ROOT/dev"

    # Run the command  in the chroot
    chroot "$ROOT" bash -c "cd $RUNDIR && eval ${COMMAND[*]}"

    if [ -n "$PRESERVE" ] ; then
        if [ "${PRESERVE:0:1}" != "/" ] ; then
            PRESERVE="$(pwd)/$PRESERVE"
        fi
        cp -a $ROOT$PRESERVE $BASE/copy
        if [ -n "$REF_UID" ] ; then
            find $BASE/copy -uid 0 | xargs chown $REF_UID.$REF_GID
        fi
        rsync -a $BASE/copy/ "$PRESERVE/"
        rm -rf $BASE/copy
    fi

    # And run the cleanup
    run_exit_trap
}

usage()
{
    cat <<EOF
usage: $me [-w WORKSPACE] [-p DIR] [-r REF-FILE] <command>
  command        - The name of an executable or script to run
  -w WORKSPACE   - The directory in which to create working directories for
                   the chroot environment.  Not required if the WORKSPACE
                   environment variable is defined.
  -p DIR         - Preserve anything in the specified directory.  Anything that
                   is changed in DIR will be merged back into the live directory
                   after the executable has completed.
  -r REF-FILE    - If -p is used to preserve a directory, because the executable
                   runs as root any new files ar owned by root.  This will change
                   the ownership of any files owner by root to be owned by the
                   same user and group as REF-FILE.

This script will run another executable in a chroot environment.  This chroot will look just like / except
any changes that are made will not persist.  The executable will run as root.  After the executable is run
a directory can be recursively merged back into the live directory, overriding the root ownership if desired.
EOF
}

# Mount an overlayfs
# $1 - The directory to mount
# $2 - The location where the overlay should be mounted
# $3 - Base directory where we can put the upperdir and workdir
mount_overlay()
{
    local LOWER="$1"
    local MOUNT="$2"
    local BASE="$3"
    local UPPER="$BASE/upper-${LOWER//\//_}"
    local WORK="$BASE/work-${LOWER//\//_}"
    mkdir -p "$UPPER" "$WORK" "$MOUNT"
    mount -t overlay -o lowerdir=${LOWER},upperdir=${UPPER},workdir=${WORK} overlay $MOUNT
    add_exit_trap "umount $MOUNT"
}

# Add a command to the existing EXIT trap
# $1 - The command to add
trapcmds=()
add_exit_trap()
{
    trapcmds=("$1" "${trapcmds[@]}")
    # On any exit, run all the commands
    trap run_exit_trap EXIT
    # On Ctrl-C, run all the commands (which disables the EXIT trap, too)
    trap "run_exit_trap ; exit 2" SIGINT
}

run_exit_trap()
{
    trap - EXIT
    for cmd in "${trapcmds[@]}" ; do
	    eval "$cmd"
    done
    trapcmds=()
}

# Retry wrapper around umount
umount()
{
    local out i
    for ((i=0; i < 20; ++i)) ; do
	if out=$(command umount "$@" 2>&1) ; then
	    #echo "umount $@"
	    return 0
	fi

	if echo "$out" | grep -q "target is busy" ; then
	    # If umount failed with "target is busy" then retry for a bit
	    sleep 0.5
	    continue
	else
	    # Failed and not sure why
	    echo "$out" >&2
	    return 1
	fi
    done
    echo "$me: error: Unable to unmount $@ after 10 attempts"
    return 1
}


main "$@"