ohmy-pentest-report.zsh-theme

# Control to show or hide the IP (disabled by default)
show_ip=false # Set to true to always enable ip

# Control to show or hide the date and time (disabled by default)
show_time=false # Set to true to always enable the date and time

# Control to use a two-line prompt (disabled by default)
twolines=false  # Set to true to enable two-line prompt

# Default interface to get the IP address (empty by default for auto-detection)
default_interface=""
interface="$default_interface"

# Variable to store a manually set IP, if needed
manual_ip=""

# Track whether public IP is being used
check_public_ip=false

# Variable to store the last time the public IP was checked
last_public_ip_check=0
public_ip_check_interval=900  # 15 minutes (900 seconds)

local user_name="%{$fg[white]%}%n"
local dir_prefix="%{$fg[red]%}["
local dir_string="%{$fg[white]%}%~"
local dir_postfix="%{$fg[red]%}] "
local cmd_symbol_success="%{$fg[white]%}❯ "  # Default prompt symbol in white
local cmd_symbol_fail="%{$fg[red]%}❯ "      # Prompt symbol in red for errors
local cmd_symbol_root="%{$fg[red]#%} "      # Symbol for root always in red
local new_line=""  # Always use a single-line prompt by default
local check_interval=10  # Check the interface every 10 seconds

# Function to get the IP address from the specified interface or the main one
get_ip_address() {
    if [ -n "$manual_ip" ]; then
        echo "$manual_ip"
    else
        local ip=""
        if [ -n "$interface" ]; then
            ip=$(ip -o -4 addr list "$interface" 2>/dev/null | awk '{print $4}' | cut -d/ -f1)
        fi
        if [ -z "$ip" ]; then
            # Get the default route interface
            local default_iface=$(ip route | awk '/default/ {print $5}' | head -n 1)
            # Get the IP address of the default route interface
            ip=$(ip -o -4 addr list "$default_iface" 2>/dev/null | awk '{print $4}' | cut -d/ -f1)
        fi
        if [ -z "$ip" ]; then
            # Get the IP address from any interface except lo
            ip=$(ip -o -4 addr list | awk '$2 != "lo" {print $4}' | cut -d/ -f1 | head -n 1)
        fi
        echo "$ip"
    fi
}

# Global variable that stores the IP in the prompt (always in magenta)
ip_prompt=""

# Function to update the IP in the prompt if it changes
update_ip_in_prompt() {
    if [ "$show_ip" = true ]; then
        local new_ip=$(get_ip_address)
        # Keep IP always in magenta
        ip_prompt="%{$fg[magenta]%}$new_ip%{$reset_color%}"
    else
        ip_prompt=""
    fi
}

# Initial update of ip_prompt
update_ip_in_prompt

# Function to update the public IP if it's been 15 minutes
update_public_ip_if_needed() {
    if [ "$show_ip" = true ] && [ "$check_public_ip" = true ]; then
        local current_time=$(date +%s)
        local time_elapsed=$((current_time - last_public_ip_check))

        if [ $time_elapsed -ge $public_ip_check_interval ]; then
            # Update the public IP
            manual_ip=$(curl -s https://ifconfig.me)
            last_public_ip_check=$current_time
            update_ip_in_prompt
        fi
    fi
}

# Function to check and update IP periodically
precmd_update() {
    if [[ $SECONDS -gt $check_interval ]]; then
        SECONDS=0
        update_ip_in_prompt
        update_public_ip_if_needed  # Check and update public IP if needed
    fi
}

# Add the precmd_update function to the precmd_functions array
typeset -ga precmd_functions
precmd_functions+=(precmd_update)

# Function to get the date and time in the desired format
get_datetime() {
    if [ "$show_time" = true ]; then
        echo "%{$fg[cyan]%}$(date +'%d/%m/%y %H:%M')%{$reset_color%}"
    else
        echo ""
    fi
}

# Configure the prompt symbol according to the user and the last command's status
set_prompt_symbol() {
    if [ "$USER" = "root" ]; then
        cmd_symbol="$cmd_symbol_root"
    else
        # Check the exit status of the previous command and set the appropriate symbol
        cmd_symbol="%(?.$cmd_symbol_success.$cmd_symbol_fail)"
    fi
}

# Unified function to set either a specific IP, obtain the IP from an interface, or get the public IP
setip() {
    if [[ -n "$1" ]]; then
        if [[ "$1" == "public" ]]; then
            # Get the public IP using an external service
            manual_ip=$(curl -s https://ifconfig.me)
            interface=""  # Clear any interface selection
            check_public_ip=true  # Enable public IP check
            last_public_ip_check=$(date +%s)  # Initialize the time of the last public IP check
        elif [[ "$1" =~ ^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+$ ]]; then
            manual_ip="$1"
            interface=""
            check_public_ip=false
        else
            interface="$1"
            manual_ip=""
            check_public_ip=false
        fi
        update_ip_in_prompt
    else
        echo "Usage: setip <ip_address|interface_name|public>"
    fi
}

# Git prompt settings (if needed)
ZSH_THEME_GIT_PROMPT_PREFIX=" %{$fg[blue]%}(%{$reset_color%}%{$fg[yellow]%}"
ZSH_THEME_GIT_PROMPT_SUFFIX="%{$reset_color%}%{$fg[blue]%})%{$reset_color%}"
ZSH_THEME_GIT_PROMPT_DIRTY="⚡"
ZSH_THEME_GIT_PROMPT_CLEAN=""

# Command to disable the IP in the prompt
disableip() {
    show_ip=false
    check_public_ip=false    # Stop checking the public IP
    manual_ip=""             # Clear any manually set IP
    interface="$default_interface"  # Reset to default interface
    update_ip_in_prompt
}

# Command to enable the IP in the prompt
enableip() {
    show_ip=true
    update_ip_in_prompt
}

# Command to disable the date and time in the prompt
disabledate() {
    show_time=false
}

# Command to enable the date and time in the prompt
enabledate() {
    show_time=true
}

# Command to enable both the IP and the date/time in the prompt
enableall() {
    show_ip=true
    show_time=true
    update_ip_in_prompt
}

# Command to disable both the IP and the date/time in the prompt
disableall() {
    show_ip=false
    show_time=false
    check_public_ip=false    # Stop checking the public IP
    manual_ip=""             # Clear any manually set IP
    interface="$default_interface"  # Reset to default interface
    update_ip_in_prompt
}

# Function to construct the prompt
construct_prompt() {
    local prompt=""
    local datetime="$(get_datetime)"
    local ip="$ip_prompt"
    local hyphen=""
    local separator=""
    local newline_char=""

    if [ -n "$datetime" ] && [ -n "$ip" ]; then
        hyphen="%{$fg[white]%} - %{$reset_color%}"
    fi

    # Include a space if datetime, hyphen, or ip is non-empty
    if [ -n "${datetime}${hyphen}${ip}" ]; then
        separator=" "
    fi

    # If twolines is true, set new_line as a newline character
    if [ "$twolines" = true ]; then
        new_line=$'\n'
    else
        new_line=""
    fi

    set_prompt_symbol  # Update the prompt symbol according to the status

    prompt="${datetime}${hyphen}${ip}${separator}${dir_prefix}${dir_string}${dir_postfix}${new_line}${cmd_symbol}%{$reset_color%}"

    echo "$prompt"
}

# Final prompt string
PROMPT='$(construct_prompt)'

# Set up initial values
update_ip_in_prompt

# Function to log commands to ~/.pentest_history
preexec_log_command() {
    local cmd=${1:-}
    if [[ -n "$cmd" ]]; then
        # Get the date in format dd/mm/yy HH:MM
        local date="$(date +'%d/%m/%y %H:%M')"
        # Get the IP address
        local ip="$(get_ip_address)"
        # Append to ~/.pentest_history
        echo "$date - $ip - $cmd" >> ~/.pentest_history
    fi
}

# Add the preexec_log_command function to the preexec_functions array
typeset -ga preexec_functions
preexec_functions+=(preexec_log_command)