#!/bin/bash
# =============================================================================
# Custom prompt configuration module
# =============================================================================

# Set script directory
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

# Source common functions and variables
source "./common.sh"

# Function to configure custom banner
configure_custom_banner() {
  log_message "INFO" "Configuring custom banner"

  # Create custom banner file
  local banner_file="/etc/banner"

  log_message "INFO" "Creating custom banner file"

  cat > "$banner_file" << 'EOF'

  _______ _                ____   __     __  __ _
 |__   __(_)              / __ \ / _|   |  \/  (_)
    | |   _ _ __  ___ ___| |  | | |_ ___| \  / |_ _ __   ___
    | |  | | '_ \/ __|___| |  | |  _|___| |\/| | | '_ \ / _ \
    | |  | | |_) \__ \   | |__| | |     | |  | | | | | |  __/
    |_|  |_| .__/|___/    \____/|_|     |_|  |_|_|_| |_|\___|
           | |
           |_|

AVERTISSEMENT : L'accès à ce système est réservé aux utilisateurs
dûment autorisés.

Toute tentative d'accès, d'accès à ce système sans autorisation ou
de maintien frauduleux dans ce système fera l'objet de poursuites
conformément à la politique de Tips-Of-Mine.

Tout utilisateur autorisé est informé et reconnaît que ses actions
peuvent être enregistrées, conservées et auditées conformément aux
chartes/politiques internes de aux chartes/politiques internes de
Tips-Of-Mine.
EOF

  chmod +x "$banner_file"

  local file="/etc/ssh/sshd_config"
  local line=`grep -n "#Banner none" $file | cut -d ":" -f 1`

  #echo $line

  # Vérification de la présence de la ligne AuthorizedKeysFile
  if [ -z "$line" ]; then
      echo "#Banner none" | tee -a $file
  else
      sed -i ''$line'c\Banner /etc/banner' $file > /dev/null || handle_error "Échec de "
  fi

  #
  line=`grep -n "#Banner none" $file | cut -d ":" -f 1` || handle_error "Échec de "

  log_message "SUCCESS" "Custom banner file created at $banner_file"

  # Source the prompt file to apply immediately
  source "$banner_file"
}

# Function to configure custom profile
configure_custom_profile() {
  log_message "INFO" "Configuring custom profile"

  # Create custom profile file
  local profile_file="/etc/profile.d/custom-profile.sh"

  log_message "INFO" "Creating custom profile file"

  cat > "$profile_file" << 'EOF'
#!/bin/bash
# Custom secure server profile
# Generated by security hardening script

# Check if the terminal supports colors
if [ -x /usr/bin/tput ] && tput setaf 1 >&/dev/null; then
  # Color definitions
  BLUE="\[\033[01;34m\]"
  GREEN="\[\033[01;32m\]"
  RED="\[\033[01;31m\]"
  YELLOW="\[\033[01;33m\]"
  PURPLE="\[\033[01;35m\]"
  CYAN="\[\033[01;36m\]"
  WHITE="\[\033[01;37m\]"
  RESET="\[\033[00m\]"
  BOLD="\[\033[01m\]"

  # Get server IP
  SERVER_IP=$(hostname -I | awk '{print $1}')

  # Define symbols based on privilege
  if [ "$(id -u)" -eq 0 ]; then
    # Root user - red prompt
    USER_COLOR=$RED
    PROMPT_SYMBOL="#"
  else
    # Regular user - green prompt
    USER_COLOR=$GREEN
    PROMPT_SYMBOL="$"
  fi

  # Set the prompt
  PS1="${BOLD}[${USER_COLOR}\u${RESET}${BOLD}@${CYAN}\h${RESET}${BOLD} ${YELLOW}\w${RESET}${BOLD}]${RESET}\\n${USER_COLOR}${PROMPT_SYMBOL}${RESET} "

  # Show additional security information for root users
  if [ "$(id -u)" -eq 0 ]; then
    # Display system information
    echo -e "\n${RED}ATTENTION: ROOT LOGIN${RESET}"
    echo -e "${YELLOW}System Info:${RESET}"
    echo -e "  ${CYAN}Hostname:${RESET} $(hostname)"
    echo -e "  ${CYAN}IP Address:${RESET} ${SERVER_IP}"
    echo -e "  ${CYAN}Kernel:${RESET} $(uname -r)"
    echo -e "  ${CYAN}Uptime:${RESET} $(uptime -p | sed 's/up //')"
    echo -e "  ${CYAN}Load:${RESET} $(cat /proc/loadavg | awk '{print $1 ", " $2 ", " $3}')"

    # Show recent failed login attempts
    FAILED_LOGINS=$(grep "Failed password" /var/log/auth.log | tail -5)
    if [ ! -z "$FAILED_LOGINS" ]; then
      echo -e "\n${YELLOW}Recent Failed Login Attempts:${RESET}"
      echo -e "${RED}$(grep "Failed password" /var/log/auth.log | tail -5)${RESET}"
    fi

    echo -e "\n${RED}THIS IS A SECURED SERVER - ALL ACTIONS ARE LOGGED${RESET}\n"
  fi
else
  # Simple prompt for terminals without color support
  PS1="[\u@\h \W]\\$ "
fi

# Set some useful aliases
alias ll='ls -la'
alias l='ls -l'
alias rm='rm -i'
alias cp='cp -i'
alias mv='mv -i'
alias grep='grep --color=auto'
EOF

  chmod +x "$profile_file"

  log_message "SUCCESS" "Custom profile file created at $profile_file"

  # Source the profile file to apply immediately
  source "$profile_file"
}

# Function to configure custom prompt
configure_custom_prompt() {
  log_message "INFO" "Configuring custom prompt"

  # Create custom prompt file
  local prompt_file="/etc/update-motd.d/00-basic"

  log_message "INFO" "Creating custom prompt file"

  cat > "$prompt_file" << 'EOF'
#!/bin/bash

# get load averages
IFS=" " read LOAD1 LOAD5 LOAD15 <<<$(awk '{ print $1,$2,$3 }' /proc/loadavg)
# get free memory
IFS=" " read USED AVAIL TOTAL <<<$(free -htm | awk '/Mem/ { print $3,$7,$2 }')
# get processes
PROCESS=$(ps -eo user=|sort|uniq -c | awk '{ print $2 " " $1 }')
PROCESS_ALL=$(echo "$PROCESS"| awk {'print $2'} | awk '{ SUM += $1} END { print SUM }')
PROCESS_ROOT=$(echo "$PROCESS" | awk '/root/ { print $2}')
PROCESS_USER=$(echo "$PROCESS" | awk '!/root/ { SUM += $2} END { print SUM }')
# get processors
PROCESSOR_NAME=$(awk -F": " '/model name/ { print $2 }' /proc/cpuinfo | head -1)
PROCESSOR_COUNT=$(grep -ioPc 'processor\t:' /proc/cpuinfo)

# colors
W="\e[0;39m"
G="\e[1;32m"
R="\e[1;31m"
dim="\e[2m"
undim="\e[0m"

echo -e "${W}System info:
$W  Hostname$dim····$undim: $W${HOSTNAME}
$W  Distro$dim······$undim: $W$(grep "PRETTY_NAME" /etc/*release | cut -d "=" -f 2- | sed 's/"//g')
$W  Kernel$dim······$undim: $W$(uname -sr)
$W  Uptime$dim······$undim: $W$(uptime -p)
$W  Load$dim········$undim: $G$LOAD1$W (1m), $G$LOAD5$W (5m), $G$LOAD15$W (15m)
$W  Processes$dim···$undim: $G$PROCESS_ROOT$W (root), $G$PROCESS_USER$W (user), $G$PROCESS_ALL$W (total)
$W  CPU$dim·········$undim: $W$PROCESSOR_NAME ($G$PROCESSOR_COUNT$W vCPU)
$W  Memory$dim······$undim: $G$USED$W used, $G$AVAIL$W avail, $G$TOTAL$W total"

# config
max_usage=90
bar_width=50

# disk usage: ignore zfs, squashfs & tmpfs
printf "\nDisk usage:\n"

while read line; do
    # get disk usage
    usage=$(echo "$line" | awk '{print $2}' | sed 's/%//')
    used_width=$((($usage*$bar_width)/100))
    # color is green if usage < max_usage, else red
    if [ "${usage}" -ge "${max_usage}" ]; then
        color=$R
    else
        color=$G
    fi
    # print green/red bar until used_width
    bar="[${color}"
    for ((i=0; i<$used_width; i++)); do
        bar+="="
    done
    # print dimmmed bar until end
    bar+="${W}${dim}"
    for ((i=$used_width; i<$bar_width; i++)); do
        bar+="·"
    done
    bar+="${undim}]"
    # print usage line & bar
    echo "${line}" | awk '{ printf("%-31s%+3s used out of %+4s\n", $1, $2, $3); }' | sed -e 's/^/  /'
    echo -e "${bar}" | sed -e 's/^/  /'
done < <(df -H -x zfs -x squashfs -x tmpfs -x devtmpfs -x overlay -x nfs -x nfs4 -x cifs --output=target,pcent,size | tail -n+2)

printf "\n"
EOF

  chmod +x "$prompt_file"

  log_message "SUCCESS" "Custom prompt file created at $prompt_file"

  # Source the prompt file to apply immediately
    source "$prompt_file"
}

# Main execution for custom prompt
configure_custom_banner
configure_custom_profile
configure_custom_prompt

log_message "SUCCESS" "Custom prompt configuration completed"