#!/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

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

############################# UTILITY FUNCTIONS ########################

NOT_EXISTING_FILE=103
NOT_ABSOLUTE_PATH=111
NO_ROOT_PRIVILEGES=110
NULL_EXCEPTION=11
WRONG_ANSWER=33

#######################################
# Check if the argument is null.
#
# Globals:
#   None
# Arguments:
#   argument ($1)    : Argument to check.
# Returns:
#   0                : If argument is not null.
#   NULL_EXCEPTION   : If argument is null.
# Output:
#   None
#######################################
function check_not_null() {
    [ -z "$1" ] && { error "Error: null argument $1"; return $NULL_EXCEPTION; }
    return 0
}

#######################################
# Redirect message to stderr.
#
# Globals:
#   None
# Arguments:
#   msg ($@): Message to print.
# Returns:
#   None
# Output:
#   Message printed to stderr.
#######################################
function echoerr() {
    echo "$@" 1>&2;
}

#######################################
# Print an error message to stderr and exit program.
#
# Globals:
#   None
# Arguments:
#   msg ($@)   : Message to print.
# Returns:
#   1          : The unique exit status printed.
# Output:
#   Message printed to stderr.
#######################################
function die() {
    error $@
    exit 1
}

#######################################
# Print an error message to stderr and exit program with a given status.
#
# Globals:
#   None
# Arguments:
#   status ($1)     : The exit status to use.
#   msg ($2-)       : Message to print.
# Returns:
#   $?              : The $status exit status.
# Output:
#   Message printed to stderr.
#######################################
function die_on_status() {
    status=$1
    shift
    error $@
    exit $status
}

#######################################
# Print an error message to stderr.
#
# Globals:
#   None
# Arguments:
#   msg ($@): Message to print.
# Returns:
#   None
# Output:
#   Message printed to stderr.
#######################################
function error() {
    echoerr -e "\033[1;31m$@\033[0m"
}

#######################################
# Print a warn message to stderr.
#
# Globals:
#   None
# Arguments:
#   msg ($@): Message to print.
# Returns:
#   None
# Output:
#   Message printed to stderr.
#######################################
function warn() {
    # $@: msg (mandatory) - str: Message to print
    echoerr -e "\033[1;33m$@\033[0m"
}

#######################################
# Print an info message to stdout.
#
# Globals:
#   None
# Arguments:
#   msg ($@): Message to print.
# Returns:
#   None
# Output:
#   Message printed to stdout.
#######################################
function info(){
    echo -e "\033[1;36m$@\033[0m"
}

function check_and_trap() {
    local sigs="${@:2:${#@}}"
    local traps="$(trap -p $sigs)"
    [[ $traps ]] && die "Attempting to overwrite existing $sigs trap: $traps"
    trap $@
}

function check_and_force_trap() {
    local sigs="${@:2:${#@}}"
    local traps="$(trap -p $sigs)"
    [[ $traps ]] && warn "Attempting to overwrite existing $sigs trap: $traps"
    trap $@
}



################################ 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 --force --recursive $mp || final_res=$?
        fi
    done
    $UMOUNT --force --recursive ${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: $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

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

EOF
}

function parse_arguments() {
    BINDINGS=()
    OPT_NO_UMOUNT=false
    OPT_RECURSIVE=false
    OPT_BIND="--bind"
    OPT_AVOID_BIND=false
    OPT_HELP=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 ;;
            -*) 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

    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 "$@"
