#!/bin/bash -e
# thanks for updating it claude.ai
# shyft
# 26DEC2024
# invoke via zsh -c "$(curl -fsSL https://kali.shyft.us)"

# Color Variables
green='\e[32m'
blue='\e[34m'
clear='\e[0m'

# Color Functions
ColorGreen() {
    echo -ne "$green$1$clear"
}
ColorBlue() {
    echo -ne "$blue$1$clear"
}

banner() {
    cat << "EOF"
 #    #                    ######                                                        
 #   #    ##   #      #    #     #  ####   ####  #####  ####  ##### #####    ##   #####  
 #  #    #  #  #      #    #     # #    # #    #   #   #        #   #    #  #  #  #    # 
 ###    #    # #      #    ######  #    # #    #   #    ####    #   #    # #    # #    # 
 #  #   ###### #      #    #     # #    # #    #   #        #   #   #####  ###### #####  
 #   #  #    # #      #    #     # #    # #    #   #   #    #   #   #   #  #    # #      
 #    # #    # ###### #    ######   ####   ####    #    ####    #   #    # #    # #      

Stage a new kali machine
EOF
}

setup_aliases_and_paths() {
    echo "Setting up aliases and paths..."
    
    # Ensure zshrc exists
    touch ~/.zshrc
    
    # Remove existing git aliases to avoid conflicts, only if they exist
    if grep -q "alias.*git" ~/.zshrc; then
        sed -i 's/\(alias.*git.*\)/#\1/' ~/.zshrc
    fi
    
    # Base aliases - using conditional addition
    if ! grep -qF "alias sudo='sudo -E'" ~/.zshrc; then
        echo "alias sudo='sudo -E'" | tee -a ~/.zshrc
    fi
    
    if ! grep -qF "alias xclip='xclip -selection clipboard'" ~/.zshrc; then
        echo "alias xclip='xclip -selection clipboard'" | tee -a ~/.zshrc
    fi
    
    if ! grep -qF "alias wormr='wormhole receive'" ~/.zshrc; then
        echo "alias wormr='wormhole receive'" | tee -a ~/.zshrc
    fi
    
    if ! grep -qF "alias worms='wormhole send'" ~/.zshrc; then
        echo "alias worms='wormhole send'" | tee -a ~/.zshrc
    fi
    
    # ZSH configurations
    if ! grep -qF "setopt no_share_history" ~/.zshrc; then
        echo "setopt no_share_history" | tee -a ~/.zshrc
    fi
    
    if ! grep -qF "unsetopt share_history" ~/.zshrc; then
        echo "unsetopt share_history" | tee -a ~/.zshrc
    fi
    
    if ! grep -qF "zstyle ':completion:*' rehash true" ~/.zshrc; then
        echo "zstyle ':completion:*' rehash true" | tee -a ~/.zshrc
    fi
    
    # PATH configurations - only add if not already present
    if ! grep -qF "export PATH=~/tools/bin:" ~/.zshrc; then
        echo "export PATH=~/tools/bin:$PATH" | tee -a ~/.zshrc
    fi
    
    # mcfly configuration - only add if not already present
    if ! grep -qF 'eval "$(mcfly init zsh)"' ~/.zshrc; then
        echo 'eval "$(mcfly init zsh)"' | tee -a ~/.zshrc
    fi
}

setup_minimal() {
    echo "Setting up minimal configuration..."
    
    # Create tools directory if it doesn't exist
    mkdir -p ~/tools/bin
    
    # Install base packages idempotently
    sudo apt update && sudo apt -y dist-upgrade
    for pkg in zsh tmux nano magic-wormhole xclip; do
        if ! dpkg -l | grep -q "^ii  $pkg "; then
            sudo apt install -y $pkg
        fi
    done
    
     # Install Rust if not already installed
    if ! command -v rustc >/dev/null 2>&1; then
        curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
    fi
    
    # Oh My Zsh installation - only if not already installed
    if [ ! -d "$HOME/.oh-my-zsh" ]; then
        RUNZSH='no' sh -c "$(curl -fsSL https://raw.github.com/ohmyzsh/ohmyzsh/master/tools/install.sh)"
    fi
    
    # Powerlevel10k theme - only if not already installed
    if [ ! -d "${ZSH_CUSTOM:-$HOME/.oh-my-zsh/custom}/themes/powerlevel10k" ]; then
        git clone --depth=1 https://github.com/romkatv/powerlevel10k.git ${ZSH_CUSTOM:-$HOME/.oh-my-zsh/custom}/themes/powerlevel10k 2>/dev/null
        sed -i 's/ZSH_THEME=".*"/ZSH_THEME="powerlevel10k\/powerlevel10k"/' ~/.zshrc
    fi
    
    # tmux configuration - only if not already setup
    if [ ! -d "$HOME/.tmux" ]; then
        cd ~
        git clone https://github.com/ShyftXero/.tmux.git 2>/dev/null
        ln -s -f .tmux/.tmux.conf
        cp -n .tmux/.tmux.conf.local . 2>/dev/null || true
    fi
    
    # nano configuration - only if not already configured
    if [ ! -f ~/.nanorc ] || ! grep -q "ShyftXero" ~/.nanorc; then
        curl -s https://raw.githubusercontent.com/scopatz/nanorc/master/install.sh | sh >/dev/null 2>&1
        curl -fsSL https://gist.github.com/ShyftXero/a3c481cce0feb57290ddd3288287fc1d/raw > ~/.nanorc
    fi
    
    # mcfly installation - only if not already installed
    if ! command -v mcfly >/dev/null 2>&1; then
        curl -LSfs https://raw.githubusercontent.com/cantino/mcfly/master/ci/install.sh | sudo sh -s -- --git cantino/mcfly
    fi
    
    # Setup aliases and paths
    setup_aliases_and_paths
    
    # Get uv_venv script - only append if not already present
    if ! grep -qF "uv_venv" ~/.zshrc; then
        curl -fsSL https://gist.githubusercontent.com/ShyftXero/e2b5c02108ce426e426f872cd3d47cb0/raw/uv_venv.sh >> ~/.zshrc
    fi
    
    # Handle cargo path if it exists
    if [ -f "$HOME/.cargo/env" ]; then
        if ! grep -qF "source ~/.cargo/env" ~/.zshrc; then
            echo "source ~/.cargo/env" | tee -a ~/.zshrc
        fi
    fi
    
    # Source zshrc only if shell is interactive
    if [[ $- == *i* ]]; then
        source ~/.zshrc
    fi
}

setup_kali() {
    echo "Setting up Kali-specific configuration..."
    
    # Run minimal setup first
    setup_minimal
    
    rmdir Documents Music Pictures Public Templates Videos
    
   
    
    # Install Kali tools idempotently
    sudo apt update
    sudo apt install -y kali-linux-default wireshark docker.io docker-compose golang openjdk-17-jdk

    if ! grep -qF "export GOPATH=\$HOME/go" ~/.zshrc; then
    echo "export GOPATH=\$HOME/go" | tee -a ~/.zshrc
    fi

    if ! grep -qF "export PATH=\$PATH:/usr/local/go/bin:\$GOPATH/bin" ~/.zshrc; then
        echo "export PATH=\$PATH:/usr/local/go/bin:\$GOPATH/bin" | tee -a ~/.zshrc
    fi

    # Set for current session too
    export GOPATH=$HOME/go
    export PATH=$PATH:/usr/local/go/bin:$GOPATH/bin
    
    
    # project discovery tools
    go install -v github.com/projectdiscovery/pdtm/cmd/pdtm@latest
    pdtm -install-all

    # Configure wireshark - only if not already configured
    if ! getcap /usr/bin/dumpcap | grep -q "cap_net_raw"; then
        sudo setcap 'CAP_NET_RAW+eip CAP_NET_ADMIN+eip' /usr/bin/dumpcap
    fi
    
    # Add user to groups if not already member
    for group in wireshark docker; do
        if ! groups | grep -q "\b$group\b"; then
            sudo usermod -a -G $group $USER
        fi
    done
    
    # Install VS Code if not already installed
    if ! command -v code >/dev/null 2>&1; then
        curl -OJL https://go.microsoft.com/fwlink/\?LinkID\=760868
        sudo dpkg -i code*.deb
        rm -f code*.deb
        
        # Install extensions
        code --install-extension ms-python.python
        code --install-extension mechatroner.rainbow-csv
        code --install-extension bibhasdn.unique-lines
        code --install-extension qcz.text-power-tools
        code --install-extension grapecity.gc-excelviewer
    fi
}

setup_extra() {
    echo "Setting up extra tools..."
    
    # Install Python packages idempotently
    pip3 install --user --upgrade requests requests-html pwntools flask pythonnet mitmproxy
    
    # Install security tools - only if not already present
    cd ~/tools
    for repo in "carlospolop/PEASS-ng" "longld/peda"; do
        dir_name=$(echo $repo | cut -d'/' -f2)
        if [ ! -d "$dir_name" ]; then
            git clone https://github.com/$repo 2>/dev/null || true
        fi
    done
    
    # Configure GDB-PEDA - only if not already configured
    if [ ! -f ~/.gdbinit ] || ! grep -qF "source ~/tools/peda/peda.py" ~/.gdbinit; then
        echo "source ~/tools/peda/peda.py" >> ~/.gdbinit
    fi
}

menu() {
    banner
    echo -ne "
    Bootstrap Menu
    $(ColorGreen '1)') Minimal Setup (zsh + tmux + nano + wormhole + mcfly + aliases)
    $(ColorGreen '2)') Kali Setup (includes minimal + security tools + VS Code)
    $(ColorGreen '3)') Extra Tools (additional security tools)
    $(ColorGreen '0)') Exit
    
    $(ColorBlue 'Choose an option:') "
    read a
    case $a in
        1) setup_minimal; menu ;;
        2) setup_kali; menu ;;
        3) setup_extra; menu ;;
        0) exit 0 ;;
        *) echo "Invalid option"; menu ;;
    esac
}

# Start the menu
menu