#!/usr/bin/env bash
###############################################################################
#                                                                             #
#                                                                             #
#                        EZ Chroot Shellscript Tool                           #
#                 Makes chroot'ing much less frustrating :)                   #
#                                                                             #
#                                                                             #
#               (C) Copyright 2020 Privex Inc. / Someguy123                   #
#            https://www.privex.io // https://peakd.com/@Someguy123           #
#                                                                             #
#                                                                             #
#                           --- Generated at ---                              #
#                       2020-SEP-30 05:44:58 AM UTC                           #
#                                                                             #
#                                                                             #
###############################################################################
#
# Install:
# 
#   sudo wget -O /usr/local/bin/ezchroot https://cdn.privex.io/extras/ezchroot.sh
#   # ALTERNATIVELY - if you don't have / don't like wget, you can use cURL
#   sudo curl -fsSL -o /usr/local/bin/ezchroot https://cdn.privex.io/extras/ezchroot.sh
#
#   sudo chmod +x /usr/local/bin/ezchroot
#
#   # Use ezchroot interactively
#   ezchroot
# 
#   # Specify a mount point and/or block device to mount
#   ezchroot /mnt/example /dev/sdc3
#   # Skip block device mounting completely.
#   ezchroot /mnt/example ""
#
#   # Specify a custom shell to use
#   ezchroot /mnt/example /dev/sdc3 "/usr/bin/zsh"
#   ezchroot /mnt/example "" "/usr/bin/qemu-arm-static /bin/bash"
#
#
###############################################################################
! [ -z ${ZSH_VERSION+x} ] && _SHTYPE=zsh || _SHTYPE=bash

get_var() { eval "printf '%s' \"\$$1\""; }
set_var() { eval "export $1=\"$2\""; }
_is_number() { grep -qE '^[0-9]+$'; }
is_number() {
    if (( $# > 0 )); then
        for n in "$@"; do
            _is_number <<< "$n"
            _ret=$?
            if (( _ret != 0 )); then return $_ret; fi
        done
    else
        _is_number
    fi
}
is_yes() {
    if [[ -z "$1" ]]; then return 1; fi
    if is_number "$1" && (( $1 > 0 )); then return 0; fi
    [[ "$1" == "y" || "$1" == "Y" || "$1" == "yes" || "$1" == "YES"  || "$1" == "true"  || "$1" == "TRUE" ]]
}
is_no() {
    if [[ -z "$1" ]]; then return 1; fi
    if is_number "$1" && (( $1 <= 0 )); then return 0; fi
    [[ "$1" == "n" || "$1" == "N" || "$1" == "no" || "$1" == "NO"  || "$1" == "false"  || "$1" == "FALSE" ]]
}

# var_intbool [var_name] (var_name2 var_name3...)
# convert the variable 'var_name' into an integer boolean (1=yes, 0=no) for use with bash/zsh (( ))
#
# e.g.
#     do_stuff='yes'
#     var_intbool do_stuff
#     echo "$do_stuff"    # output: 1
#
var_intbool() {
    local v_name v_val
    for v_name in "$@"; do
        v_val="$(get_var "$v_name")"
        is_yes "$v_val" && set_var "$v_name" 1 || set_var "$v_name" 0
        v_name=""
    done
}

: ${ASKQ_APPEND_DEFAULT=1}

# askq [output_variable] [question/prompt]
# askq [output_variable] [question/prompt] (default/fallback_value) (append_default_prompt=yes)
askq() {
    local out_var="$1" print_q="$2" ans_default="_NO_DEFAULT_SET_" append_prompt="$ASKQ_APPEND_DEFAULT"
    
    (( $# > 2 )) && ans_default="$3"
    (( $# > 3 )) && append_prompt="$4"
    is_yes "$append_prompt" && append_prompt=1 || append_prompt=0

    [[ "$ans_default" != "_NO_DEFAULT_SET_" ]] && (( append_prompt )) && print_q="$print_q ( default: '$ans_default' ): "
    ! [ -z ${ZSH_VERSION+x} ] && _SHTYPE=zsh || _SHTYPE=bash

    if [[ "$_SHTYPE" == "zsh" ]]; then
        vared -p "$print_q" -c "$out_var"
    else
        read -p "$print_q" "$out_var"
    fi
    if [[ -z "$(get_var $out_var)" ]] && [[ "$ans_default" != "_NO_DEFAULT_SET_" ]]; then
        echo -e "No answer. Using default value: $ans_default"
        set_var "$out_var" "$ans_default"
    fi
}

# yesno (prompt="Are you sure?") (default)
yesno() {
    local y_prompt="Are you sure?" y_default="_NO_DEFAULT_SET_" y_opts="(y/n)" y_val
    (( $# > 0 )) && y_prompt="$1"
    if (( $# > 1 )); then
        y_default="$2"
        is_yes "$y_default" && y_opts="(Y/n)"
        is_no "$y_default" && y_opts="(y/N)"
        if ! is_yes "$y_default" && ! is_no "$y_default"; then
            >&2 echo -e "\n [!!!] Invalid default '${y_default}' provided to yesno(). Not using a default.\n"
            y_default="_NO_DEFAULT_SET_" y_opts="(y/n)"
        fi
    fi
    askq y_val "$y_prompt $y_opts > " "$y_default" no
    
    if [[ -z "$y_val" ]]; then
        if [[ "$y_default" == "_NO_DEFAULT_SET_" ]]; then
            >&2 echo -e "\n [!!!] Empty reply, but this question has no default option. Please enter one of: y, Y, n, N, yes, no, YES, or NO.\n"
            yesno "$y_prompt"
            return $?
         fi
         >&2 echo -e "\n >>> No answer given. Using default answer: ${y_default}\n"
    fi
    if is_yes "$y_val"; then return 0; fi
    if is_no "$y_val"; then return 1; fi
    >&2 echo -e "\n [!!!] Invalid answer '${y_val}'. Please enter one of: y, Y, n, N, yes, no, YES, or NO.\n"
    yesno "$y_prompt" "$y_default"
    return $?
}

# yesno_intbool [var_name] (prompt="Are you sure?") (default)
yesno_intbool() {
    local v_name="$1"
    yesno "${@:2}" && set_var "$v_name" 1 || set_var "$v_name" 0
}

uses_sysd_resolve() {
    local rootfs_dir="$1" _res
    local resconf="${rootfs_dir}/etc/resolv.conf" sysd_resolvedir="${rootfs_dir}/run/systemd/resolve"
    if [[ -d "$sysd_resolvedir" ]]; then return 0; fi
    if [[ -f "${rootfs_dir}/etc/systemd/resolved.conf" ]]; then return 0; fi
    if grep -q "nameserver 127.0.0." "$resconf"; then return 0; fi
    if grep -qi "systemd-resolved" "$resconf"; then return 0; fi
    if grep -qi "stub" "$resconf"; then return 0; fi
    return 1
}

_inj_ns() {
    local pdir="$(dirname "$1")"
    [[ ! -d "$pdir" ]] && echo -en " >> Parent folder $pdir not found. Creating it... " && mkdir -pv "$pdir"
    touch "$1"
    echo "nameserver 185.130.44.20" >> "$1"
    echo "nameserver 1.1.1.1" >> "$1"
    echo "nameserver 2a07:e00::333" >> "$1"
}

fix_resolve() {
    local rootfs_dir="${CHRDIR}"
    (( $# > 0 )) && rootfs_dir="$1"
    local res_dir="${rootfs_dir}/run/systemd/resolve"
    local res_stub="${res_dir}/stub-resolv.conf" resolv_conf="${rootfs_dir}/etc/resolv.conf"

    if uses_sysd_resolve "$rootfs_dir"; then
        echo -e " >> System at '${rootfs_dir}' appears to use systemd-resolved. Initialising ${res_stub} ..."
        _inj_ns "$res_stub"
    else
        echo -e " >> System at '${rootfs_dir}' doesn't appear to use systemd-resolved. Writing directly to $resolv_conf ..."
        _inj_ns "$resolv_conf"
    fi
}

bind_mounts() {
    if [[ ! -d proc ]] || [[ -z "$(ls proc)" ]]; then
        echo -e " [...] ${PWD}/proc not mounted. Mounting now."
        echo -e "         -> $(mount -v -t proc proc proc/) \n"
    else
        echo -e " [+++] ${PWD}/proc is non-empty - must be already mounted. Skipping.\n"
    fi
    if [[ ! -d sys ]] || [[ -z "$(ls sys)" ]]; then
        echo -e " [...] ${PWD}/sys not mounted. Mounting now."
        echo -e "         -> $(mount -v --bind /sys sys/) \n"
    else
        echo -e " [+++] ${PWD}/sys is non-empty - must be already mounted. Skipping.\n"
    fi
    if [[ ! -d dev ]] || [[ -z "$(ls dev)" ]] || ! grep -q "${PWD}/dev" <<< "$(mount)"; then
        echo -e " [...] ${PWD}/dev not mounted. Mounting now."
        echo -e "         -> $(mount -v --bind /dev dev/) \n"
    else
        echo -e " [+++] ${PWD}/dev is non-empty - must be already mounted. Skipping.\n"
    fi
}


: ${INITFILE_DEF_LOC="var/tmp/.initfile"}

_mk-initfile() {
    local init_out="${CHRDIR}/${INITFILE_DEF_LOC}"
    # (( $# > 0 )) && init_out="$1"
    {
        echo '[[ -f /etc/bash.bashrc ]] && source /etc/bash.bashrc'
        echo '[[ -f "${HOME}/.bashrc" ]] && source "${HOME}/.bashrc"'
        echo 'export PATH="/bin:/sbin:/usr/sbin:/usr/bin:/usr/local/bin:/usr/local/sbin:$PATH"'
        echo 'hostname "$(cat /etc/hostname)"'
        echo 'export HOSTNAME="$(cat /etc/hostname)"'
        is_yes "$CHR_MOUNT_FSTAB" && echo 'echo -e "   >>> Mounting all filesystems from /etc/fstab now ..."' && echo 'mount -av'
        echo 'TERM=xterm-256color' "$CHR_SHELL_BASE"
        echo
    } > "$init_out"
    chmod +x "$init_out" &>/dev/null
    echo "/$INITFILE_DEF_LOC"
}

: ${AUTO_SET_NS=0}
: ${AUTO_MTAB=0}
: ${AUTO_DEV=0}
: ${AUTO_CHRDIR=0}
: ${AUTO_SHELL=0}
ezchroot() {
    : ${CHR_SHELL_BASE="/bin/bash"}
    : ${CHR_SHELL="/bin/bash"}
    : ${CHR_DEV=""}
    : ${CHRDIR="/mnt/root"}
    : ${CHR_SET_NS='yes'}
    : ${CHR_MTAB='yes'}
    : ${CHR_MOUNT_FSTAB='no'}
    : ${CHR_INIT_FILE=1}
    local cr_dev init_file_loc=""
    
    (( $# > 0 )) && CHRDIR="$1" && AUTO_CHRDIR=1
    (( $# > 1 )) && cr_dev="$2" && AUTO_DEV=1
    (( $# > 2 )) && CHR_SHELL="$3" && AUTO_SHELL=1

    is_yes "$CHR_MTAB" && CHR_MTAB='yes' || CHR_MTAB='no'
    is_yes "$CHR_SET_NS" && CHR_SET_NS='yes' || CHR_SET_NS='no'

    echo -e "If you haven't yet mounted the partition/volume you want to chroot into, you may enter the path to"
    echo -e "it's block device / image file (e.g. /dev/vda2) at the following question and we'll mount it for you.\n"

    echo -e "If you don't need us to mount the chroot base filesystem for you, just hit enter at the prompt.\n"

    (( AUTO_DEV )) || askq cr_dev "(optional) Block device / image file to mount: " "$CHR_DEV"
    echo
    if [[ -n "$cr_dev" ]]; then
        echo -e " >> You've selected the block device '${cr_dev}'. At the following question, you should enter"
        echo -e " >> the mountpoint folder where you'd like to mount it. It will be auto-created if it doesn't already exist.\n"
    fi
    (( AUTO_CHRDIR )) || askq CHRDIR "Enter the path to the filesystem you want to chroot into" "$CHRDIR"
    echo
    if [[ -n "$cr_dev" ]]; then
        [[ ! -d "$CHRDIR" ]] && echo -en " [+++] Folder '${CHRDIR}' not found. Creating it..." && echo -e " [---]" "$(mkdir -pv "$CHRDIR")"
        grep -q "$CHRDIR" <<< "$(mount)" || echo -en "\n [+++] Device/Image '${cr_dev} not mounted at $CHRDIR - mounting now..." && \
                echo -e " [---]" "$(mount -v "$cr_dev" "$CHRDIR")" "\n"
    fi
    (( AUTO_SHELL )) || askq CHR_SHELL "Shell to use within the chroot (inc. arguments if needed)" "$CHR_SHELL"
    (( AUTO_SHELL )) || askq CHR_SHELL_BASE "Shell / interpreter to be launched by the init script (inside the chroot) after the chroot is activated" "$CHR_SHELL_BASE"
    # if grep -q '/bash' <<< "$CHR_SHELL"; then
    yesno_intbool CHR_MOUNT_FSTAB "Auto-mount all filesystems in ${CHRDIR}/etc/fstab after entering the chroot?" "$CHR_MOUNT_FSTAB"
    # fi
    echo
    (( AUTO_MTAB )) || yesno_intbool CHR_MTAB "Auto-copy /etc/mtab from host system if needed?" "$CHR_MTAB"
    echo
    (( AUTO_SET_NS )) || yesno_intbool CHR_SET_NS "Auto-configure DNS resolver configs in the chroot ( e.g. /etc/resolv.conf ) ?" "$CHR_SET_NS"
    
    var_intbool CHR_SET_NS CHR_MTAB
    is_yes "$CHR_SET_NS" && CHR_SET_NS=1 || CHR_SET_NS=0
    echo

    cd "$CHRDIR"
    bind_mounts
    (( CHR_SET_NS )) && fix_resolve "$CHRDIR"
    IFS=" " read -r -a CHR_SHELL <<< "$CHR_SHELL"
    (( CHR_INIT_FILE )) && init_file_loc="$(_mk-initfile)"

    chroot_args=( "$CHRDIR" "${CHR_SHELL[@]}" )
    if [[ -n "$init_file_loc" ]]; then
        chroot_args+=("$init_file_loc")
    fi
    
    echo -e "\n !!!               Entering CHROOT now                 !!!\n"
    echo -e "      Command:\n"
    echo -e "          chroot ${chroot_args[*]}"
    echo -e "\n !!!               Entering CHROOT now                 !!!\n"

    chroot "${chroot_args[@]}"
}



ezchroot "${@:1}"