#!/bin/bash
#
# This script is the simplified and portable version of arch-chroot
# (https://wiki.archlinux.org/index.php/Change_root#Using_arch-chroot)
#

set -e

# JUNEST_BASE can be overridden for testing purposes.
# There is no need for doing it for normal usage.
JUNEST_BASE="${JUNEST_BASE:-$(readlink -f $(dirname $(readlink -f "$0"))/..)}"
NAME='GRoot'
CMD='groot'
DESCRIPTION="I am $NAME!"
CHROOTCMD=${CHROOTCMD:-chroot}
SHELL="/bin/sh"
MOUNT=mount
UMOUNT=umount
MKDIR=mkdir
TOUCH=touch
CUT=cut
SORT=sort
UNIQ=uniq
CAT=cat
READLINK=readlink
MOUNTS_FILE=/proc/self/mounts

NOT_EXISTING_FILE=103
NOT_ABSOLUTE_PATH=111
NO_ROOT_PRIVILEGES=110

source "${JUNEST_BASE}/lib/utils/utils.sh"


################################ MAIN FUNCTIONS ###########################

function is_mountpoint() {
    local mountpoint="$1"
    for mp in $($CAT $MOUNTS_FILE | $CUT -f2 -d' ' | $SORT -r | $UNIQ)
    do
        [[ $mp == $mountpoint ]] && return 0
    done
    return 1
}

function chroot_teardown() {
    # Remove all mounts starting from the most nested ones.
    # Suffix the CHROOTDIR with / to avoid umounting directories not belonging
    # to CHROOTDIR.
    local normalized_chrootdir="$($READLINK -f ${CHROOTDIR})/"
    local final_res=0
    for mp in $($CAT $MOUNTS_FILE | $CUT -f2 -d' ' | $SORT -r | $UNIQ)
    do
        if [[ $mp =~ ^${normalized_chrootdir}.* ]] && is_mountpoint "$mp"
        then
            $UMOUNT $mp || final_res=$?
        fi
    done
    $UMOUNT ${CHROOTDIR%/}

    return $final_res
}

function chroot_maybe_add_mount() {
    local cond=$1
    shift
    if eval "$cond"; then
        $MOUNT "$@"
        return
    fi
    return 1
}

function chroot_setup() {
    $OPT_NO_UMOUNT || check_and_trap 'chroot_teardown' QUIT EXIT ABRT KILL TERM INT

    if ! chroot_maybe_add_mount "! is_mountpoint '$CHROOTDIR'" --bind "$CHROOTDIR" "$CHROOTDIR"
    then
        warn "Failed mount of directories. $CHROOTDIR is already a mountpoint. Skipping it..."
        return 0
    fi

    local re='(.*):(.*)'
    for binds in ${BINDINGS[@]}
    do
        local host_path=""
        local guest_path=""
        if [[ $binds =~ $re ]]
        then
            local host_path="${BASH_REMATCH[1]}"
            local guest_path="${BASH_REMATCH[2]}"
        else
            local host_path="$binds"
            local guest_path="$binds"
        fi

        create_node "${host_path}" "${CHROOTDIR}${guest_path}"
        mount_directory "${host_path}" "${guest_path}"
    done
}

function mount_directory() {
    local host_path=$($READLINK -f "$1")
    local guest_path="$2"

    if ! $OPT_AVOID_BIND
    then
        $MOUNT $OPT_BIND "${host_path}" "${CHROOTDIR}${guest_path}"
        return 0
    fi

    case "$host_path" in
        /proc) $MOUNT proc "${CHROOTDIR}${guest_path}" -t proc ;;
        /sys) $MOUNT sys "${CHROOTDIR}${guest_path}" -t sysfs ;;
        /dev) $MOUNT udev "${CHROOTDIR}${guest_path}" -t devtmpfs; $MOUNT devpts "${guest_path}/pts" -t devpts; $MOUNT shm "${guest_path}/shm" -t tmpfs ;;
        /run) $MOUNT run "${CHROOTDIR}${guest_path}" -t tmpfs ;;
        /tmp) $MOUNT tmp "${CHROOTDIR}${guest_path}" -t tmpfs ;;
        *) $MOUNT $OPT_BIND "${host_path}" "${CHROOTDIR}${guest_path}" ;;
    esac

    return 0
}

function create_node() {
    local src="$1"
    local dst="$2"
    if [[ ! -e $src ]]
    then
        die_on_status $NOT_EXISTING_FILE "${src} does not exist."
    elif [[ $src != /* ]]
    then
        die_on_status $NOT_ABSOLUTE_PATH "${src} is not an absolute path."
    elif [[ -f $src ]]
    then
        $TOUCH "$dst"
    elif [[ -d $src ]]
    then
        $MKDIR -p "$dst"
    fi
}

function usage() {
  cat <<EOF
$NAME (v$(cat $JUNEST_BASE/VERSION)): $DESCRIPTION

Usage:
    $CMD [options] [<chroot-dir> [command]]

Options:
    -b, --bind <path>
            Make the content of <path> accessible in the guest rootfs.

            This option makes any file or directory of the host rootfs
            accessible in the confined environment just as if it were part of
            the guest rootfs.  By default the host path is bound to the same
            path in the guest rootfs but users can specify any other location
            with the syntax: -b <host_path>:<guest_location>. This option can
            be invoked multiple times and the paths specified must be absolutes.

    -n, --no-umount
            Do not umount after chroot session finished.

    -r, --recursive
            Use rbind instead of bind.

    -i, --avoid-bind
            Attempt to avoid mount --bind for common directories and use
            proper mount fstype instead. Detected directories with
            corresponding fstype are: /proc (proc), /sys (sysfs),
            /dev (devtmpfs), /tmp (tmpfs), /run (tmpfs).

    -h, --help                Print this help message

    -V, --version             Show the $NAME version

If 'command' is unspecified, $CMD will launch $SHELL.

EOF
}

version() {
    echo -e "$NAME $(cat $JUNEST_BASE/VERSION)"
}

function parse_arguments() {
    BINDINGS=()
    OPT_NO_UMOUNT=false
    OPT_RECURSIVE=false
    OPT_BIND="--bind"
    OPT_AVOID_BIND=false
    OPT_HELP=false
    OPT_VERSION=false
    for opt in "$@"
    do
        case "$1" in
            -b|--bind) shift ; BINDINGS+=("$1") ; shift ;;
            -n|--no-umount) OPT_NO_UMOUNT=true ; shift ;;
            -r|--recursive) OPT_BIND="--rbind" ; shift ;;
            -i|--avoid-bind) OPT_AVOID_BIND=true ; shift ;;
            -h|--help) OPT_HELP=true ; shift ;;
            -V|--version) OPT_VERSION=true ; shift ;;
            -*) die "Invalid option $1" ;;
            *) break ;;
        esac
    done

    if [[ ! -z $1 ]]
    then
        CHROOTDIR="$1"
        shift
    fi
    ARGS=()
    for arg in "$@"
    do
        ARGS+=("$arg")
    done
}

function is_user_root() {
    (( EUID == 0 ))
}

function execute_operation() {
    $OPT_HELP && usage && return
    $OPT_VERSION && version && return

    is_user_root || die_on_status $NO_ROOT_PRIVILEGES 'This script must be run with root privileges'

    [[ -d $CHROOTDIR ]] || die_on_status $NOT_EXISTING_FILE "Can't create chroot on non-directory $CHROOTDIR"

    chroot_setup "$CHROOTDIR" || die "Failed to setup chroot $CHROOTDIR"

    $CHROOTCMD "$CHROOTDIR" "${ARGS[@]}"
}


function main() {
    parse_arguments "$@"
    execute_operation
}

main "$@"