#! /bin/bash
#  Copyright 2023 NVIDIA Corporation
#  Copyright 2016 by Cumulus Networks, Inc.

# This script needs to be run as root or via sudo

# Set up an existing account to be use a restricted shell
# (only rbash is supported), changing the home directory so the
# user can not remove the dot files, and have the dot files
# owned by bin.bin.

# The intent is to set up an account for use by users who
# are only allowed to use commands (and command arguments)
# permitted by a TACACS+ server.

# Set up a ~user/bin directory and PATH to point to that.
# Create a symlink from /usr/sbin/tacplus-auth to each of
# the the requested commands to the bin directory

# Those will be the only commands accessible to the user, and
# only if also permitted by the TACACS+ server

# Use the -i option to initialize the directory
# permissions and dot files.  It will only be run
# once, unless the -f option is given

# Use the -a to add command links that the user will be able
# to potentially run (subject to authorization by the TACACS+ server)
# The list of commands is given after all other arguments

# -i and -a can be combined

# The -u username argument is required

# See
#   man tacplus-auth
# for additional information.

prog=${0##*/}

# could add rzsh or rksh in future, but keeping it simple for now
shell=/bin/rbash origshell=/bin/bash
tacacs_restrict=.tacacs_restricted # flag file for inited

authprog=/usr/sbin/tacplus-auth # command that does authorization

# set defaults
user= firstarg=
addcmds=0 delcmds=0 doinit=0 dorestore=0 force=0 serverauth=0 noserverauth=0

warn()
{
    echo $prog: Warning: "$@" 1>&2
}

usage()
{
   echo Usage: $prog '-u user [-f (force)] [-i (init)] [-A (server auth)] [-D (disable server auth)] [-R (restore)] [(-a|-d) cmd1 cmd2 ...]' 1>&2
   exit 1
}

fatal()
{
   warn "$@"
   exit 2
}

# Generate the .bashrc configuration for restricted bash with history and NVUE password censoring
generate_rbash_config()
{
    cat << 'EOF'
if [ "$0" = "-rbash" ]; then
    # Enable history with recommended settings
    shopt -s histappend
    export HISTCONTROL=ignoreboth
    export HISTSIZE=1000
    export HISTFILESIZE=2000
    
    # Keep tab completion enabled (commented out the restriction)
    # shopt -u progcomp
    
    # ===== NVUE builtins-only history censor for ultra-restricted rbash =====
    # No redirections, no external tools, no brace groups. Builtins only.
    
    __nvue_censor_hook() {
      local saved_flags=$-
      set -f  # disable globbing
    
      # Get last history entry (number + command)
      local last
      last=$(history 1)
      if [[ -z "$last" ]]; then
        [[ "$saved_flags" != *f* ]] && set +f
        return
      fi
    
      # Normalize whitespace via array split/join
      local _arr=() norm
      read -ra _arr <<< "$last"
      norm="${_arr[*]}"
    
      # Split history number and command
      norm="${norm#"${norm%%[![:space:]]*}"}"
      local hist_num cmd
      hist_num="${norm%% *}"
      cmd="${norm#* }"
      if [[ -z "$cmd" ]]; then
        [[ "$saved_flags" != *f* ]] && set +f
        return
      fi
    
      # Tokenize command
      local t=()
      read -ra t <<< "$cmd"
    
      # Heuristic rules
      local next_kws=(secret password pass token key psk community passphrase)
      local opt_kws=(--secret --password --pass --token --key --psk --passphrase)
    
      local out=()
      local i=0
      local in_pem=0
    
      while (( i < ${#t[@]} )); do
        local w="${t[i]}"
    
        # Collapse PEM blocks
        if (( in_pem )); then
          if [[ "$w" == -----END* ]]; then
            in_pem=0
          fi
          (( i++ ))
          continue
        fi
        if [[ "$w" == -----BEGIN* ]]; then
          out+=( "*" )
          in_pem=1
          (( i++ ))
          continue
        fi
    
        # Handle --opt=value
        local matched=0
        local k
        for k in "${opt_kws[@]}"; do
          if [[ "$w" == "$k="* ]]; then
            out+=( "${k}=*" )
            matched=1
            break
          fi
        done
        if (( matched )); then
          (( i++ ))
          continue
        fi
    
        # Handle --opt value
        matched=0
        for k in "${opt_kws[@]}"; do
          if [[ "$w" == "$k" ]]; then
            out+=( "$w" )
            (( i++ ))
            if (( i < ${#t[@]} )); then
              out+=( "*" )
              (( i++ ))
            fi
            matched=1
            break
          fi
        done
        if (( matched )); then
          continue
        fi
    
        # Handle bare keyword value
        matched=0
        for k in "${next_kws[@]}"; do
          if [[ "$w" == "$k" ]]; then
            out+=( "$w" )
            (( i++ ))
            if (( i < ${#t[@]} )); then
              out+=( "*" )
              (( i++ ))
            fi
            matched=1
            break
          fi
        done
        if (( matched )); then
          continue
        fi
    
        # Mask long token-ish strings
        if [[ "${#w}" -ge 32 && "$w" =~ ^[A-Za-z0-9+/=._-]+$ ]]; then
          out+=( "*" )
          (( i++ ))
          continue
        fi
    
        # Default passthrough
        out+=( "$w" )
        (( i++ ))
      done
    
      # Write back to history (builtins only)
      if [[ -n "$hist_num" ]]; then
        history -d "$hist_num" || true
      fi
      if (( ${#out[@]} > 0 )); then
        history -s "${out[@]}"
      fi
    
      [[ "$saved_flags" != *f* ]] && set +f
    }
    
    # IMPORTANT: do NOT "source /path/..." in rbash; call the function by name.
    PROMPT_COMMAND="__nvue_censor_hook"
    # ===== end hook =====
fi
EOF
}

# Generate the _upvars function override for bash completion
generate_upvars_override()
{
    cat << 'EOF'
# Override _upvars function from bash_completion to avoid stderr redirection issues
_upvars()
{
    if ! (($#)); then
        echo "bash_completion: $FUNCNAME: usage: $FUNCNAME" \
             "[-v varname value] | [-aN varname [value ...]] ..."
        return 2
    fi
    while (($#)); do
        case $1 in
            -a*)
                [[ ${1#-a} ]] || {
                    echo "bash_completion: $FUNCNAME:" \
                         "\`$1': missing number specifier"
                    return 1
                }
                # FIXED: Regex check instead of redirection
                [[ "${1#-a}" =~ ^[0-9]+$ ]] || {
                    echo bash_completion: \
                         "$FUNCNAME: \`$1': invalid number specifier"
                    return 1
                }
                [[ "$2" ]] && unset -v "$2" && eval $2=\(\"\$"{@:3:${1#-a}}"\"\) &&
                    shift $((${1#-a} + 2)) || {
                    echo bash_completion: \
                         "$FUNCNAME: \`$1${2+ }$2': missing argument(s)"
                    return 1
                }
                ;;
            -v)
                [[ "$2" ]] && unset -v "$2" && eval $2=\"\$3\" &&
                    shift 3 || {
                    echo "bash_completion: $FUNCNAME: $1:" \
                         "missing argument(s)"
                    return 1
                }
                ;;
            *)
                echo "bash_completion: $FUNCNAME: $1: invalid option"
                return 1
                ;;
        esac
    done
}
EOF
}

# Generate the _filedir function override for bash completion
generate_filedir_override()
{
    cat << 'EOF'
# Override _filedir function from bash_completion to avoid stderr redirection issues
_filedir()
{
    local i IFS=$'\n' xspec
    _tilde "$1" || return 0
    local -a toks
    local quoted x tmp
    _quote_readline_by_ref "$1" quoted
    # FIX: Removed '2>/dev/null' from compgen
    x=$( compgen -d -- "$quoted" ) &&
    while read -r tmp; do
        toks+=( "$tmp" )
    done <<< "$x"
    if [[ "$1" != -d ]]; then
        # Munge xspec to contain forbidden chars
        [[ ${1:-} ]] && xspec="${1:+"!*.$1"}" || xspec=
        
        # FIX: Removed '2>/dev/null' from compgen
        x=$( compgen -f -X "$xspec" -- "$quoted" ) &&
        while read -r tmp; do
            toks+=( "$tmp" )
        done <<< "$x"
    fi
    if [[ ${#toks[@]} -ne 0 ]]; then
        # FIX: Removed '2>/dev/null' from compopt
        compopt -o filenames
        COMPREPLY+=( "${toks[@]}" )
    fi
}
EOF
}

initialize()
{
    set -e # exit if any of the next group of commands fail

    chsh -s $shell $user || fatal Use tacacs0, tacacs15, etc. for -u arg
    chmod 755 $home || $warn Unable to chmod ${user} home directory $home

    dots=$(echo .[a-zA-Z]*)
    if [ -n "${dots}" -a "${dots}" != ".[a-zA-Z]*" ]; then
	    # create dir if needed.  If this is the 2nd init
	    # or later with dotfiles, the original backup will
	    # be saved in the new backup.
	    mkdir -p DOTback
	    mv -i .[a-zA-Z]* DOTback/
	    chown bin:bin DOTback
	    chmod 750 DOTback
	    mv DOTback .DOTback
	    echo Moved old dot files to DOTback: ; (cd .DOTback && ls -a)
    fi

    for f in .bash_login .profile .bash_profile .bash_logout
    do echo '. ~/.bashrc' > $f
    done

    echo "export PATH=$home/bin" > .bashrc

    # Add rbash-specific configuration with history and NVUE password censoring
    # (see generate_rbash_config function for details)
    generate_rbash_config >> .bashrc

    # Add bash completion function overrides
    # These override the standard bash_completion functions to avoid stderr redirection
    # which is not allowed in restricted bash
    generate_upvars_override >> .bashrc
    generate_filedir_override >> .bashrc

    mkdir -p bin
    touch ${tacacs_restrict}
    chown -R bin:bin .
    chmod -R o-w . # be paranoid

    mkdir -p .local
    chown -R $user:tacacs .local
    chmod -R u+rwX .local

    # Create .bash_history file with proper ownership so rbash can write to it on exit
    touch .bash_history
    chown $user:tacacs .bash_history
    chmod u+rw .bash_history

    set +e # initialization complete
}


# dorestore the initialization.  Return shell to /bin/bash, restore dotfiles, etc.
restore()
{
    local l cmd lev=${user#tacacs}

    warn Restoring $user '(tacacs priv='$lev')' to normal use,'
    able to run all commands.'
    warn You have 5 seconds to interrupt to prevent the restore
    sleep 5

    chsh -s $origshell $user || warn unable to restore shell $origshell for $user

    for f in .bash_login .profile .bash_profile .bash_logout
    do rm -f $f # remove the files we created in initialize()
    done

    dots=$(echo .DOTback/.[a-zA-Z]* | sed 's,.DOTback/.DOTback,,')
    if [ -z "${dots}" -o "${dots}" = ".DOTback/.[a-zA-Z]*" ]; then
        warn No saved files '(.bashrc, etc.)', using /etc/skel
        cp -a /etc/skel/.[a-zA-Z0-9]* .
    else
        mv ${dots} .
    fi

    for f in . .??* ; do
        chown tacacs${lev}:tacacs $f
    done
    chown -R tacacs${lev}:tacacs .local

    [ -e ${tacacs_restrict} ] && rm ${tacacs_restrict}

    # Only remove .DOTback if empty, not rm -rf.  Sometimes multiple levels
    rmdir .DOTback 2>/dev/null
    [ -d .DOTback ] && chown -R ${user}:tacacs .DOTback # make user can see files

    for cmd in bin/*; do
        [ -L "${cmd}" ] && {
            l=$(readlink -s "${cmd}")
            case "$l" in
            $authprog) rm -f "${cmd}" ;;
            esac
        }
    done
    rmdir bin 2>/dev/null # remove if empty
}

# Enable server-side per-command authorization (-A flag).
# Creates the tacplus-auth symlink and appends command_not_found_handle
# to .bashrc so that ANY command the user types is forwarded to the
# TACACS+ server for authorization, without needing per-command symlinks.
enable_serverauth()
{
    [ -x $authprog ] || fatal TACACS authorization command $authprog not found
    [ -e ${tacacs_restrict} ] || warn not yet initialized, run with -i first

    mkdir -p bin
    ln -s -f $authprog bin/tacplus-auth
    echo Server-side authorization symlink created: bin/tacplus-auth

    # Append command_not_found_handle to .bashrc if not already present
    if grep -q 'command_not_found_handle' .bashrc 2>/dev/null; then
        warn command_not_found_handle already present in .bashrc, skipping
    else
        cat >> .bashrc << 'SERVERAUTH'

# ===== command_not_found_handle for server-side per-command authorization =====
# When a command is not found in PATH (i.e., no symlink exists in ~/bin),
# forward it to tacplus-auth for TACACS+ server-side authorization.
# This allows the TACACS+ server to be the sole authority on which
# commands are permitted, without needing per-command symlinks.
# Enabled via: tacplus-restrict -u <user> -A
if [ "$0" = "-rbash" ]; then
    command_not_found_handle() {
        if [ "$1" = "tacplus-auth" ]; then
            echo "tacplus-auth: command not found (check symlink)" >&2
            return 127
        fi
        tacplus-auth "$@"
        return $?
    }
fi
# ===== end command_not_found_handle =====
SERVERAUTH
        chown bin:bin .bashrc
        chmod 444 .bashrc
        echo Server-side command authorization enabled in .bashrc
    fi
}

# Disable server-side per-command authorization (-D flag).
# Removes the tacplus-auth symlink and strips the command_not_found_handle
# block from .bashrc, reverting to symlink-only command authorization.
disable_serverauth()
{
    # Remove the tacplus-auth symlink if it points to authprog
    if [ -L bin/tacplus-auth ]; then
        l=$(readlink -s bin/tacplus-auth)
        case "$l" in
        $authprog) rm -f bin/tacplus-auth
                   echo Server-side authorization symlink removed: bin/tacplus-auth ;;
        *) warn bin/tacplus-auth points to "$l", not $authprog, skipping ;;
        esac
    else
        warn bin/tacplus-auth symlink not found, skipping
    fi

    # Remove the command_not_found_handle block from .bashrc
    if grep -q 'command_not_found_handle' .bashrc 2>/dev/null; then
        # Remove lines between the marker comments (inclusive)
        sed -i '/^# ===== command_not_found_handle for server-side/,/^# ===== end command_not_found_handle/d' .bashrc
        chown bin:bin .bashrc
        chmod 444 .bashrc
        echo Server-side command authorization removed from .bashrc
    else
        warn command_not_found_handle not found in .bashrc, skipping
    fi
}

linkcmds()
{
    [ -x $authprog ] || fatal TACACS authorization command $authprog not found
    for c in "${firstarg}" "$@"; do
       cmd=${c##*/} # make relative in case they gave full name
       case $cmd in
       sh|dash|bash|csh|tcsh|zsh|ksh|sh.distrib)
	  echo $prog: Skipping shell $cmd because it bypasses restrictions
          continue ;;
       esac
       ln -s -f $authprog bin/$cmd
    done
}

# inverse of linkcmds, but supports '*' to delete all
unlinkcmds()
{
    (
    [ ! -d bin ] && { warn All commands for $user have already been removed; exit 1; }
    cd bin || exit 1
    for c in "${firstarg}" "$@"; do
       [ -L "${c}" ] || warn "${c}" is not currently an enabled command # but remove anyway
       case "${c}" in
       '*') warn Removing all commands for $user: $(echo *) ; rm -f -- * ;;
       *)  rm -f -- "${c}" ;;
       esac
    done
    )
}

listcmds()
{
    local lev=${user#tacacs}
    if [ ! -d bin ]; then
        warn No commands available to $user, tacacs priv=${lev}.
        return
    fi
    (
    cd bin || exit 1
    p=${user/tacacs}
    case "$p" in
    [0-9]|1[0-5]) umsg="$user (TACACS+ priv=$p)" ;;
    *) umsg=$user ;;
    esac
    cmds=$(echo *)
    case "$cmds" in
    '*'|'') warn No commands available to user $umsg ;;
    *)      echo Commands available to $umsg are:
            /bin/ls -C
    esac
    )
}

# With both arrays and using set -- $(getopt ...), if any of the
# getopts has similar issues, and I want to detect things like
#   -a cmd1 cmd2 -d cmd3
# getopts builtin can't do that, but getopt does.
# There is an issue with quoted arguments not remaining quoted
# but that's OK for this use.
# Unfortunately, the quoting itself (getopt uses apostrophe quotes)
# can cause problems.  So use getopts, and step through all args.

while getopts a:d:fhiADRu: opt
do
    case "$opt" in
    a) addcmds=1 firstarg="$OPTARG" ;;
    d) delcmds=1 firstarg="$OPTARG" ;;
    u) user=$OPTARG ;;
    i) doinit=1 ;;
    f) force=1 ;;
    A) serverauth=1 ;;
    D) noserverauth=1 ;;
    R) dorestore=1;;
    h|?) usage ;;
    esac
done

shift $(( OPTIND - 1 ))

# catch things like -a b c -d e f g
for a in $@; do
    case "$a" in
    -d) [ $addcmds -eq 1 ] && usage ;;
    -a) [ $delcmds -eq 1 ] && usage ;;
    esac
done

# sanity check that we have valid argument sets, and no
# extra arguments
[ $addcmds -eq 1 -a $delcmds -eq 1 ] && usage
[ $serverauth -eq 1 -a $noserverauth -eq 1 ] && usage
[ $addcmds -eq 0 -a $delcmds -eq 0 -a $# -ne 0 ] && usage
[ -z "$user" ] && usage
[ $doinit -eq 1 -a $addcmds -eq 0 -a $delcmds -eq 0 -a $# -ne 0 ] && usage
[ $addcmds -eq 0 -a $doinit -eq 0 -a $delcmds -eq 0 -a $dorestore -eq 0 -a $serverauth -eq 0 -a $noserverauth -eq 0 ] && usage
[ $dorestore -eq 1 -a \( $addcmds -eq 1 -o $delcmds -eq 1 -o $doinit -eq 1 -o $serverauth -eq 1 -o $noserverauth -eq 1 -o $# -ne 0 \) ] && usage

home=$(eval echo ~"${user}")
[ -z "$home" -o "$home" = ~${user} ] &&
	fatal Unable to get home directory for $user

cd $home || fatal home dir "${home}" for $user does not exist

if [ -f ${tacacs_restrict} -a $doinit -eq 1 -a $force -eq 0 ]; then
   warn $user already initialized and restricted
   doinit=0
fi

[ \( $addcmds -eq 1 -o $delcmds -eq 1 \) -a ! -e ${tacacs_restrict} ] &&
	warn not yet initialized

[ -x $shell ] || fatal shell "${shell}" is not executable

umask 222 # no write perms on any dirs or files

[ $dorestore -eq 1 ] && { restore ; exit 1; }

[ $doinit -eq 1 ] && initialize

[ $serverauth -eq 1 ] && enable_serverauth
[ $noserverauth -eq 1 ] && disable_serverauth

[ $addcmds -eq 1 ] && linkcmds "$@"
[ $delcmds -eq 1 ] && unlinkcmds "$@"

listcmds # Always show which commands are avilable now
