#!/bin/bash
AXIOM_PATH="$HOME/.axiom"
if [ -f "$AXIOM_PATH/interact/includes/vars.sh" ] ; then
source "$AXIOM_PATH/interact/includes/vars.sh"
fi
cat << EOF

 ▄▄▄      ▒██   ██▒ ██▓ ▒█████   ███▄ ▄███▓
▒████▄    ▒▒ █ █ ▒░▓██▒▒██▒  ██▒▓██▒▀█▀ ██▒
▒██  ▀█▄  ░░  █   ░▒██▒▒██░  ██▒▓██    ▓██░
░██▄▄▄▄██  ░ █ █ ▒ ░██░▒██   ██░▒██    ▒██
 ▓█   ▓██▒▒██▒ ▒██▒░██░░ ████▓▒░▒██▒   ░██▒
 ▒▒   ▓▒█░▒▒ ░ ░▓ ░░▓  ░ ▒░▒░▒░ ░ ▒░   ░  ░
  ▒   ▒▒ ░░░   ░▒ ░ ▒ ░  ░ ▒ ▒░ ░  ░      ░
  ░   ▒    ░    ░   ▒ ░░ ░ ░ ▒  ░      ░
      ░  ░ ░    ░   ░      ░ ░         ░

The Dynamic Infrastructure Framework for Everybody - @pry0cc & @0xtavian

We are lucky enough to be sponsored by the awesome SecurityTrails! 
Sign up for your free account here! https://securitytrails.com/app/account?utm_source=axiom

EOF

###########################################################################################################
# Shell setup functions
#
function bash_shell(){
    echo -e "${Blue}Backing up $(echo "$HOME"/.bashrc) to $(echo "$HOME"/.bashrcbak) just in case.${Color_Off}"
    cp "$HOME"/.bashrc "$HOME"/.bashrcbak >> /dev/null 2>&1
    if [[ -d "/usr/local/go/" ]] ; then
    echo -e "${Blue}Golang already installed, skipping...${Color_Off}" 
    else
    echo -e "${Blue}Installing Golang 1.19.1${Color_Off}"
    wget -q https://golang.org/dl/go1.19.1.linux-amd64.tar.gz && sudo tar -C /usr/local -xzf go1.19.1.linux-amd64.tar.gz && rm go1.19.1.linux-amd64.tar.gz
    fi
    mkdir -p "${HOME}/go"
    export GOPATH=$HOME/go
    export PATH=$GOPATH/bin:/usr/local/go/bin:$PATH
    echo "export GOPATH=\$HOME/go" >> "${HOME}"/.bashrc
    echo "export PATH=\$GOPATH/bin:/usr/local/go/bin:\$PATH:\$HOME/.local/bin" >> "${HOME}"/.bashrc
    source "${HOME}"/.bashrc
    SHELL=$(which bash)
    echo -e "${Green}You're running Bash! Installing Axiom to \$PATH...${Color_Off}"
    echo "export PATH=\"\$PATH:\$HOME/.axiom/interact\"" >>~/.bashrc
    echo "[[ -f ~/.bashrc ]] && . ~/.bashrc" >> "${HOME}"/.bash_profile
    source ~/.bashrc  >> /dev/null 2>&1
}

function zsh_shell(){
    sudo apt install zsh -y
    echo -e "${Blue}Backing up $(echo "$HOME"/.zshrc) to $(echo "$HOME"/.zshrcbak) just in case.${Color_Off}"
    cp "$HOME"/.zshrc "$HOME"/.zshrcbak >> /dev/null 2>&1
    if [[ -d "/usr/local/go/" ]] ; then
    echo -e "${Blue}Golang already installed, skipping...${Color_Off}" 
    else
    echo -e "${Blue}Installing Golang 1.19.1${Color_Off}"
    wget -q https://golang.org/dl/go1.19.1.linux-amd64.tar.gz && sudo tar -C /usr/local -xzf go1.19.1.linux-amd64.tar.gz && rm go1.19.1.linux-amd64.tar.gz
    fi
    mkdir -p "${HOME}/go"
    export GOPATH=$HOME/go
    export PATH=$GOPATH/bin:/usr/local/go/bin:$PATH
    echo "export GOPATH=\$HOME/go" >>~/.zshrc
    echo "export PATH=\$GOPATH/bin:/usr/local/go/bin:\$PATH:\$HOME/.local/bin"  >>~/.zshrc
    echo -e "${Green}You're running ZSH! Installing Axiom to \$PATH...${Color_Off}"
    echo "export PATH=\"\$PATH:\$HOME/.axiom/interact\"" >>~/.zshrc  
    echo "source $HOME/.axiom/functions/autocomplete.zsh" >>~/.zshrc 
    source "${HOME}"/.zshrc >> /dev/null 2>&1
    SHELL=$(which zsh)
}

function omz_shell(){
    echo -e "${Blue}Backing up $(echo "$HOME"/.zshrc) to $(echo "$HOME"/.zshrcbak) just in case.${Color_Off}"
    cp "$HOME"/.zshrc "$HOME"/.zshrcbak >> /dev/null 2>&1
    sudo apt install zsh zsh-syntax-highlighting zsh-autosuggestions -y
    sh -c "$(curl -fsSL https://raw.githubusercontent.com/ohmyzsh/ohmyzsh/master/tools/install.sh)" "" --unattended
    if [[ -d "/usr/local/go/" ]] ; then
    echo -e "${Blue}Golang already installed, skipping...${Color_Off}" 
    else
    echo -e "${Blue}Installing Golang 1.19.1${Color_Off}"
    wget -q https://golang.org/dl/go1.19.1.linux-amd64.tar.gz && sudo tar -C /usr/local -xzf go1.19.1.linux-amd64.tar.gz && rm go1.19.1.linux-amd64.tar.gz
    fi
    mkdir -p "${HOME}/go"
    export GOPATH=$HOME/go
    export PATH=$GOPATH/bin:/usr/local/go/bin:$PATH
    echo "export GOPATH=\$HOME/go" >>~/.zshrc
    echo "export PATH=\$GOPATH/bin:/usr/local/go/bin:\$PATH:\$HOME/.local/bin"  >>~/.zshrc
    echo "export PATH=\$GOPATH/bin:/usr/local/go/bin:\$PATH" >>~/.zshrc
    echo -e "${Green}You're running ZSH! Installing Axiom to \$PATH...${Color_Off}"
    echo "export PATH=\"\$PATH:\$HOME/.axiom/interact\"" >>~/.zshrc  
    echo "source $HOME/.axiom/functions/autocomplete.zsh" >>~/.zshrc   
    source "${HOME}"/.zshrc  >> /dev/null 2>&1
    SHELL=$(which zsh)
}

###########################################################################################################
# Help Menu
# 
function usage() {
        echo -e "${BWhite}Description:"
        echo -e "${BWhite}Examples:${Color_Off}"
        echo -e "  ${Blue}axiom-configure${Color_Off} # kick off axiom configuration wizard"
        echo -e "  ${Blue}axiom-configure --shell Zsh${Color_Off} # Set default shell to Zsh via command line"
        echo -e "  ${Blue}axiom-configure --config <axiom.json config data>${Color_Off} # Use plaintext axiom.json data for account setup TODO: support axiom.json as a file"
        echo -e "  ${Blue}axiom-configure --shell BASH --unattended --config <axiom.json config data> ${Color_Off} # Use bash for default shell, do not promot, use plaintext axiom.json data for account setup"
        echo -e "${BWhite}Usage:${Color_Off}"
        echo -e "  --config <axiom.json data>"
        echo -e "    Setup account with data supplied from the command line (must be vaild data from axiom.json wrapped in single quotes)"
        echo -e "  --unattended (optional)"
        echo -e "    Avoid any promots, defaults to Oh My Zsh unless --shell argument is also provided"
        echo -e "  --shell <Bash, Zsh, OMZ (Oh My Zsh)> (optional)"
        echo -e "    Choose bash/zsh to add axiom to .bashrc/.zshrc. OMZ installs Oh My Zsh and overwrites .zshrc"
        echo -e "  --help"
        echo -e "    Display this help menu"
}

###########################################################################################################
# Parse command line arguments 
#
i=0
for arg in "$@"
do
    i=$((i+1))
    if [[  ! " ${pass[@]} " =~ " ${i} " ]]; then
        set=false
        if [[ "$arg" == "--config" ]]; then
            n=$((i+1))
            config=$(echo ${!n})
            config_flag="true"
            set=true
            pass+=($i)
            pass+=($n)
        fi
        if [[ "$arg" == "--unattended" ]]; then
            unattended=true          
            set=true
            pass+=($i)
        fi
        if [[ "$arg" == "--shell" ]]; then
            n=$((i+1))
            usershell=$(echo ${!n})
            set=true
            pass+=($i)
            pass+=($n)
        fi
        if [[ "$arg" == "--help" ]] || [[ "$arg" == "-h" ]] || [[ "$arg" == "help" ]]; then
            usage
            exit
            set=true
            pass+=($i)
        fi
fi
done

if [ ! -z ${unattended:+x} ]; then
echo -e "${Blue}Hello! Welcome to Axiom configuration :)"
echo -e "${Green}Detected unattended option passed from command line...${Color_Off}"
echo -e "${Blue}Automatically installing dependencies and begining installation!${Color_Off}"
if [ -z ${usershell:+x} ]; then
usershell="omz"
fi
elif [ ! -z ${config+x} ]; then
echo -e "${Blue}Hello! Welcome to Axiom configuration :)"
echo -e "${Green}Detected axiom.json option passed from command line...${Color_Off}"
echo -e "${Blue}Automatically installing dependencies and begining installation!${Color_Off}"
else
echo -e "${Blue}Hello! Welcome to Axiom configuration :)"
echo -e "Press enter to install dependencies and begin installation!${Color_Off}"
read
echo -e "${Blue}INSTALLING AXIOM...${Color_Off}"
sleep 1
fi

echo -e "${BWhite}Installing deps! Please wait :) ${Color_Off}"
BASEOS="$(uname)"
case $BASEOS in
'Linux')
    BASEOS='Linux'
    ;;
'FreeBSD')
    BASEOS='FreeBSD'
    alias ls='ls -G'
    ;;
'WindowsNT')
    BASEOS='Windows'
    ;;
'Darwin')
    BASEOS='Mac'
    ;;
'SunOS')
    BASEOS='Solaris'
    ;;
'AIX') ;;
*) ;;
esac

if [[ $BASEOS == "Linux" ]]; then
    if $(uname -a | grep -qi "Microsoft"); then
        OS="UbuntuWSL"
    else
        OS=$(lsb_release -i | awk '{ print $3 }')
        if ! command -v lsb_release &> /dev/null; then
            echo "WARNING: Unless using Ubuntu latest, this install might not work"
            echo "lsb_release could not be found, unable to determine your distribution"
            OS="unknown-Linux"
            BASEOS="Linux"
        fi
    fi
    if [[ $OS == "Arch" ]] || [[ $OS == "ManjaroLinux" ]]; then
        echo -e "${Blue}Congrats, you run arch..."
        echo -e "${Blue}Installing other repo deps...${Color_Off}"
        # Stated dependencies in the Github-wiki
        sudo pacman -Syu git ruby curl jq packer rsync --noconfirm
        # Other dependencies
        sudo pacman -Syu go python-pip net-tools unzip libxslt bc --noconfirm
        echo -e "${Blue}Installing Interlace...${Color_Off}"
        sudo rm -fr /tmp/interlace
        git clone https://github.com/codingo/Interlace.git /tmp/interlace
        cd /tmp/interlace && sudo python3 setup.py install
        
    elif [[ $OS == "Ubuntu" ]] || [[ $OS == "Debian" ]] || [[ $OS == "Linuxmint" ]] || [[ $OS == "Parrot" ]] || [[ $OS == "Kali" ]] || [[ $OS == "unknown-Linux" ]] || [[ $OS == "UbuntuWSL" ]]; then
       if ! command -v sudo &> /dev/null; then
        export DEBIAN_FRONTEND=noninteractive
        apt-get update && apt-get install sudo curl jq tzdata -y -qq
        ip=$(curl -s https://ifconfig.me/)
        ln -fs /usr/share/zoneinfo/$(curl ipinfo.io/"$ip" | jq -r .timezone) /etc/localtime
        dpkg-reconfigure --frontend noninteractive tzdata
       fi
        echo -e "${Blue}Installing other repo deps...${Color_Off}"
        DEBIAN_FRONTEND=noninteractive sudo apt-get update && sudo apt-get install git ruby python3-pip curl iputils-ping net-tools unzip xsltproc bc rsync sudo wget nano bsdmainutils openssh-server fail2ban -y 
        echo -e "${Blue}Installing jq...${Color_Off}"
        sudo wget -q -O /usr/bin/jq https://github.com/stedolan/jq/releases/download/jq-1.6/jq-linux64 && sudo chmod +x /usr/bin/jq
        echo -e "${Blue}Installing packer...${Color_Off}"
        wget -q -O /tmp/packer.zip https://releases.hashicorp.com/packer/1.8.1/packer_1.8.1_linux_amd64.zip && cd /tmp/ && unzip packer.zip && sudo mv packer /usr/bin/ && rm /tmp/packer.zip 
        echo -e "${Blue}Installing Interlace...${Color_Off}"
        sudo rm -fr /tmp/interlace
        git clone https://github.com/codingo/Interlace.git /tmp/interlace
        cd /tmp/interlace && sudo python3 setup.py install

        elif [[ $OS == "Fedora" ]]; then
        echo -e "${Blue}Installing other repo deps...${Color_Off}"
        sudo dnf update && sudo dnf -y install mmv bc fzf git rubypick python3-pip curl net-tools unzip util-linux fail2ban
        echo -e "${Blue}Installing jq...${Color_Off}"
        sudo wget -q -O /usr/bin/jq https://github.com/stedolan/jq/releases/download/jq-1.6/jq-linux64 && sudo chmod +x /usr/bin/jq
        echo -e "${Blue}Installing packer...${Color_Off}"
        wget -O /tmp/packer.zip https://releases.hashicorp.com/packer/1.8.1/packer_1.8.1_linux_amd64.zip && cd /tmp/ && unzip packer.zip && sudo mv packer /usr/bin/ && rm /tmp/packer.zip
        fi
fi

if [[ $BASEOS == "Mac" ]]; then
         echo -e "${Blue}Installing Brew...${Color_Off}"
         /bin/bash -c "$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/HEAD/install.sh)"
         echo -e 'Setting permissions'
         sudo chown -R $(whoami) /usr/local/Homebrew
         sudo chown -R $(whoami) /usr/local/var/homebrew
         sudo chown -R $(whoami) /usr/local/etc/bash_completion.d /usr/local/lib/pkgconfig /usr/local/share/aclocal /usr/local/share/doc /usr/local/share/info /usr/local/share/locale /usr/local/share/man /usr/local/share/man/man1 /usr/local/share/man/man3 /usr/local/share/man/man5 /usr/local/share/man/man7
         sudo chown -R $(whoami) /usr/local/share/zsh /usr/local/share/zsh/site-functions
         chmod u+w /usr/local/share/zsh /usr/local/share/zsh/site-functions
         chmod u+w /usr/local/etc/bash_completion.d /usr/local/lib/pkgconfig /usr/local/share/aclocal /usr/local/share/doc /usr/local/share/info /usr/local/share/locale /usr/local/share/man /usr/local/share/man/man1 /usr/local/share/man/man3 /usr/local/share/man/man5 /usr/local/share/man/man7
         echo -e "${Blue}Installing wget...${Color_Off}"
         brew install wget
         echo -e "${Blue}Installing go...${Color_Off}"
         brew install golang
         mkdir -p "${HOME}/go"
         export GOPATH=$HOME/go
         echo "export GOPATH=\$HOME/go" >>~/.zshrc
         echo "export PATH=\$GOPATH/bin:/usr/local/go/bin:\$PATH:\$HOME/.local/bin"  >>~/.zshrc
         echo -e "${Green}You're running ZSH! Installing Axiom to \$PATH...${Color_Off}"
         echo "export PATH=\"\$PATH:\$HOME/.axiom/interact\"" >>~/.zshrc
         echo -e "${Blue}Installing jq...${Color_Off}"
         brew install jq
         echo -e "${Blue}Installing coreutils...${Color_Off}"
         brew install coreutils
         echo -e "${Blue}Installing Python3...${Color_Off}"
         brew install python3
         echo -e "${Blue}Installing Interlace...${Color_Off}"
         sudo rm -fr /tmp/interlace
         git clone https://github.com/codingo/Interlace.git /tmp/interlace
         cd /tmp/interlace && sudo python3 setup.py install
         echo -e "${Blue}Installing packer...${Color_Off}"
         brew tap hashicorp/tap
         brew install hashicorp/tap/packer
         brew upgrade hashicorp/tap/packer
         source "${HOME}"/.zshrc
         SHELL=$(which zsh)
fi

if [[ $BASEOS == "Linux" ]] ; then
if [[ $usershell == "omz" ]]|| [[ $usershell == 'Oh my zsh' ]] || [[ $usershell == 'Oh My Zsh' ]] || [[ $usershell == "OMZ" ]] ; then
omz_shell
elif [[ $usershell == "bash" ]]|| [[ $usershell == "Bash" ]] || [[ $usershell == "BASH" ]] ; then
bash_shell
elif [[ $usershell == "zsh" ]]|| [[ $usershell == "Zsh" ]] || [[ $usershell == "ZSH" ]] ; then
zsh_shell
else

# Pick your shell
echo -e "${Blue}Choose bash/zsh to add axiom to .bashrc/.zshrc. OMZ installs Oh My Zsh and overwrites .zshrc${Color_Off}"
 PS3="Please select an option : "
  choices=("Bash" "Zsh" 'Oh My Zsh')
  select choice in "${choices[@]}"; do
   case $choice in
   Bash)
   bash_shell
  break
   ;;
   Zsh)
   zsh_shell
  break
   ;;
  'Oh My Zsh')
   omz_shell
  break
   ;;
esac
done
fi
fi


###########################################################################################################
# Clone axiom repo
#
if [ ! -d "$AXIOM_PATH" ]; then
    echo -e "${Blue}Installing axiom scripts...${Color_Off}"
    git clone https://github.com/pry0cc/axiom "$AXIOM_PATH"
else
    cd "$AXIOM_PATH" && git pull
fi

###########################################################################################################
#  Install anew if not already installed
#
if ! [ -x "$(command -v anew)" ]; then
go install github.com/tomnomnom/anew@latest >> /dev/null 2>&1
fi

###########################################################################################################
# unattended setup flag logic 
#
if [ -z ${config:+x} ]; then
if [[ "$unattended" == "true" ]]; then
echo "Unattended setup"
mkdir  -p "$HOME"/.axiom/accounts/

else
###########################################################################################################
# Default setup logic ( e.g. when --unattended and --config args are not used
#
mkdir -p "${HOME}/.axiom/tmp/"
"$AXIOM_PATH/interact/account-helpers/generate_amass.sh"
echo -e "${Blue}Configuring your cloud account...${Color_Off}"
"$AXIOM_PATH"/interact/axiom-account-setup


###########################################################################################################
# If axiom_rsa doesnt exist, create the SSH key and add to it Packer Builder and axiom.json
#
account_path=$(ls -la "$AXIOM_PATH"/axiom.json | rev | cut -d " " -f 1 | rev)
if [ -f ~/.ssh/axiom_rsa ] ; then
 echo -e "${Blue}Adding axiom_rsa key to Packer builder${Color_Off}"
 jq '.sshkey="'axiom_rsa'"' <"$account_path">"$AXIOM_PATH"/tmp.json ; mv "$AXIOM_PATH"/tmp.json "$account_path"
 cat ~/.ssh/axiom_rsa.pub >"$AXIOM_PATH"/configs/authorized_keys

else
 echo -e "${Blue}Creating axiom SSH key in ~/.ssh/axiom_rsa.pub${Color_Off}"
 ssh-keygen -b 2048 -t rsa -f ~/.ssh/axiom_rsa -q -N ""
 jq '.sshkey="'axiom_rsa'"' <"$account_path">"$AXIOM_PATH"/tmp.json ; mv "$AXIOM_PATH"/tmp.json "$account_path"  >> /dev/null 2>&1
 cat ~/.ssh/axiom_rsa.pub >"$AXIOM_PATH"/configs/authorized_keys
fi
fi 

else
###########################################################################################################
# Config parsing logic
#
echo "Parsing account.json file passed from command line"
provider="$(echo "$config" | jq -r '.provider')"
mkdir  -p "$HOME"/.axiom/accounts/
account=${provider}$(date +%s)
echo "$config" > "$HOME"/.axiom/accounts/"$account".json
echo "$config" > "$HOME"/.axiom/axiom.json
echo -e "${Blue}Cloud account configuration passed from command line...${Color_Off}"
fi

export PATH="$PATH:$HOME/.axiom/interact"

if [[ $unattended == "true" ]] ; then
echo -e "${BWhite}Axiom install complete. Run axiom-account-setup to use axiom...${Color_Off}"
echo -e "${BWhite}You may need to restart your shell: $SHELL${Color_Off}"
if [[ $config_flag == "true" ]] ; then
"$HOME"/.axiom/interact/axiom-account "$account"  
fi
exit 0
fi

if [ -z ${config+x} ]; then
echo -e "${BWhite}Beginning first packer build...${Color_Off}"
cd "$AXIOM_PATH"/ || exit
"$AXIOM_PATH"/interact/axiom-build
cd || exit
echo -e "${BWhite}You may need to restart your shell: $SHELL${Color_Off}"
else
echo -e "${BWhite}You may need to restart your shell: $SHELL${Color_Off}"
"$HOME"/.axiom/interact/axiom-account "$account"
fi