#!/bin/bash

AWS_REGION=us-east-1
EC2_NAME=pentestec2
EC2_KEY_NAME=ec2-pentest-key
EC2_SECURITY_GROUP_NAME=ec2-pentest-usage
MANAGEMENT_SCRIPT=ec2-utils.sh

install_prerequisities() {
    if [[ ! -z "$(which ruby)" ]] && [[ ! -z "$(which gem)" ]] && [[ ! -z "$(which bundle)" ]]; then
        echo "[.] Prerequisities already installed."
        return
    fi

    sudo=sudo
    if [ $EUID -eq 0 ] || [[ ! -z "$(which sudo)" ]]; then
        sudo=
    fi

    if [[ ! -z "$(which dnf)" ]]; then
        $sudo dnf update
        $sudo dnf install -y ssh cron ruby rubygems awscli jq
    elif [[ ! -z "$(which apt-get)" ]]; then
        $sudo apt-get update
        $sudo apt-get install -y ssh cron ruby rubygems awscli jq
    elif [[ ! -z "$(which yum)" ]]; then
        $sudo yum update
        $sudo yum install -y ssh cron ruby rubygems awscli jq
    else
        echo "[!] Please install following packages: ssh, cron, ruby, rubygems, awscli, jq yourself first."
        exit 1
    fi

    gem install bundler
    bundle install

    aw=$(which aws)
    if [ -z "$aw" ]; then
        echo "[!] Something went wrong. There is no 'aws' after installing 'awscli'."
        exit 1
    fi
}

configure_credentials() {
    if [ -e ~/.aws/credentials ]; then
        con=$(cat ~/.aws/credentials)
        if echo $con | grep -q AKIA; then
            echo "[.] AWS CLI already configured."
            return
        fi
    fi

    if [ ! -d "~/.aws" ]; then
        mkdir -p ~/.aws
    fi

    access_key=
    secret_key=
    read -p 'AWS Access Key (AKIA...): ' access_key
    read -s -p "AWS Secret Access Key: " secret_key

    if [ "$secret_key" == "" ] || [ "$access_key" == "" ]; then
        echo "[!] You must specify both access key and secret key."
        exit 1
    fi

    touch ~/.aws/config
    touch ~/.aws/credentials

    cat <<EOT > ~/.aws/config
[default]
region = $AWS_REGION
output = json
EOT

    cat <<EOT > ~/.aws/credentials
[default]
aws_access_key_id = $access_key
aws_secret_access_key = $secret_key
EOT
}

create_security_group() {
    groupId=$(aws ec2 create-security-group --group-name $EC2_SECURITY_GROUP_NAME --description "Security Group with allow any any for inbound and outbound." 2>&1 )

    if echo $groupId | grep -q "InvalidGroup.Duplicate"; then
        echo "[.] Security Group already exists."
        return
    elif echo $groupId | grep -q "An error occurred" ; then
        echo "[!] Failed: $groupId"
        exit 1
    fi

    groupId=$(echo $groupId | jq -r .GroupId)

    echo "[>] Security Group ID: $groupId"
    aws ec2 authorize-security-group-ingress --group-id $groupId --ip-permissions IpProtocol=tcp,FromPort=0,ToPort=65535,IpRanges='[{CidrIp=0.0.0.0/0,Description="allow any any"}]'
    aws ec2 authorize-security-group-ingress --group-id $groupId --ip-permissions IpProtocol=udp,FromPort=0,ToPort=65535,IpRanges='[{CidrIp=0.0.0.0/0,Description="allow any any"}]'
    aws ec2 authorize-security-group-egress --group-id $groupId --ip-permissions IpProtocol=tcp,FromPort=0,ToPort=65535,IpRanges='[{CidrIp=0.0.0.0/0,Description="allow any any"}]'
    aws ec2 authorize-security-group-egress --group-id $groupId --ip-permissions IpProtocol=udp,FromPort=0,ToPort=65535,IpRanges='[{CidrIp=0.0.0.0/0,Description="allow any any"}]'

    echo "[>] Created inbound/outbound allow any any rules"
}

delete_key_pair_and_instances() {
    echo "[.] Deleting key pair..."
    out=$(aws ec2 delete-key-pair --key-name $EC2_KEY_NAME 2>&1 ) 
    if echo $out | grep -q "An error occurred" ; then
        echo "[!] Deleting key pair failed: $out"
        exit 1
    fi

    echo "[.] Terminating related EC2 instances..."
    out=$(aws ec2 terminate-instances --instance-ids $(aws ec2 describe-instances --query 'Reservations[].Instances[].InstanceId' --filters "Name=key-name,Values=$EC2_KEY_NAME" --output text) 2>&1 ) 
    if echo $out | grep -q "An error occurred" ; then
        echo "[!] Terminating instances failed: $out"
        exit 1
    fi
}

create_ec2_key_pair() {
    keymaterial=
    out=$(aws ec2 describe-key-pairs --key-name $EC2_KEY_NAME 2>&1 ) 
    if echo $out | jq -r .KeyPairs[0].KeyName | grep -q $EC2_KEY_NAME ; then
        echo "[.] Key pair already created in your AWS account."

        filefinger=
        if [ -f "~/.aws/$EC2_KEY_NAME.pem" ] ; then
            filefinger=$(openssl pkcs8 -in ~/.aws/$EC2_KEY_NAME.pem -inform PEM -outform DER -topk8 -nocrypt | openssl sha1 -c | cut -d= -f2 | tr -d ' ')
        fi
        awsfinger=$(echo $out | jq -r .KeyPairs[0].KeyFingerprint)
        
        if [ ! -e ~/.aws/$EC2_KEY_NAME.pem ]; then
            echo
            echo "[!] ERROR: It looks like you don't have ~/.aws/$EC2_KEY_NAME.pem file with your EC2 Key-Pair"
            echo "[!] This means you do have the key pair on your AWS account but not physically in a file."
            echo "[!] In such case you will be unable to SSH into previously created EC2 instances without it, so "
            echo "[!] you are left with creating a new key pair and then terminating existing EC2 instances."
            echo
            read -p "Do you want to remove your $EC2_NAME EC2 instance and recreate key pairs? [y/N] " -n 1 -r
            echo
            if [[ $REPLY =~ ^[Yy]$ ]]
            then
                delete_key_pair_and_instances
            else
                echo "[-] Unable to continue. You must sort this out yourself."
                return
            fi
        elif [ "$filefinger" != "" ] && [ "$awsfinger" != "" ] && [ "$filefinger" != "$awsfinger" ]; then
            echo
            echo "[!] ERROR: It looks like your EC2 Key Pair located at ~/.aws/$EC2_KEY_NAME.pem"
            echo "[!] has differrent fingerprint than the Key Pair on your AWS account."
            echo "[!] In such case you will be unable to SSH into previously created EC2 instances using it, so "
            echo "[!] you are left with creating a new key pair and then terminating existing EC2 instances."
            echo
            echo "Your locally stored PEM fingerprint: $filefinger"
            echo "Your on AWS stored PEM fingerprint: $awsfinger"
            echo
            read -p "Do you want to remove your $EC2_NAME EC2 instance and recreate key pairs? [y/N] " -n 1 -r
            echo
            if [[ $REPLY =~ ^[Yy]$ ]]
            then
                delete_key_pair_and_instances
            else
                echo "[-] Unable to continue. You must sort this out yourself."
                return
            fi
        else
            echo "[+] EC2 Key pair already created. PEM stored at: ~/.aws/$EC2_KEY_NAME.pem"
            return
        fi
    fi

    keymaterial=$(aws ec2 create-key-pair --key-name $EC2_KEY_NAME 2>&1 ) 
    if echo $keymaterial | grep -q "An error occurred" ; then
        echo "[!] Creating key pair failed: $keymaterial"
        exit 1
    fi

    touch ~/.aws/$EC2_KEY_NAME.pem
    echo $keymaterial | jq -r .KeyMaterial > ~/.aws/$EC2_KEY_NAME.pem
    chmod 400 ~/.aws/$EC2_KEY_NAME.pem
    
    echo "[+] EC2 Key pair created. PEM stored at: ~/.aws/$EC2_KEY_NAME.pem"
}

integrate_with_bashrc() {
    if cat ~/.bashrc | grep -q $MANAGEMENT_SCRIPT ; then
        echo "[.] Script is already integrated."
        return
    elif cat ~/.bashrc | grep -q "Adding aliases for EC2 single-instance management" ; then
        echo "[.] Script is already integrated."
        return
    fi

    sed -i -e "s@ruby ./aws-manager.rb@ruby $PWD/aws-manager.rb@" ./$MANAGEMENT_SCRIPT

    source $PWD/$MANAGEMENT_SCRIPT
    cat <<EOT >> ~/.bashrc

# Adding aliases for EC2 single-instance management.
source $PWD/$MANAGEMENT_SCRIPT

EOT
}

add_cron_job() {
    crontab -l > /tmp/.cr 2> /dev/null
    cp /tmp/.cr /tmp/.cr.bak

    if cat /tmp/.cr | grep -q "aws-manager.rb notify "; then
        echo "[.] Crontab job already scheduled."
        return
    fi

    cmd="$(which ruby) $PWD/aws-manager.rb notify $EC2_NAME"
    job="  *  */2 *  *  * $USER    $cmd"

    echo -e "[*] Adding following line to crontab:\n\n\t\t$job\n"

    echo $job >> /tmp/.cr
    crontab /tmp/.cr
    if [ $? -ne 0 ]; then
        echo "[!] Updating crontab failed. Restoring original one (you can find it at: /tmp/.cr.bak) "
        crontab /tmp/.cr.bak
    fi
    rm /tmp/.cr
    #rm /tmp/.cr.bak
}

echo
echo "----------------------------------------------"
echo
echo ":: AWS EC2 single-instance management utilities installation script."
echo
echo "This script is going to:"
echo -e "\t- Update your repos & install packages such as: ssh, cron, jq, ruby, rubygems, awscli, gem bundler, gem 'aws-sdk-ec2'"
echo -e "\t- Configure your AWS credentials"
echo -e "\t- Create AWS security groups, EC2 key pairs"
echo -e "\t- Integrate EC2 management aliases into the end of your .bashrc"
echo -e "\t- Add a cron job that will notify you every two hours if your EC2 machine is up and running"
echo
echo "----------------------------------------------"
echo
read -p "Would you like to proceed? [Y/n] " -n 1 -r
echo
if [[ $REPLY =~ ^[Yy]$ ]]
then
    echo -e "\n[.] Step 1: Installing prerequisities first."
    install_prerequisities

    echo -e "\n[.] Step 2: Configuring credentials."
    configure_credentials

    echo -e "\n\n[.] Step 3: Creating Security Group."
    create_security_group

    echo -e "\n[.] Step 4: Creating EC2 Key Pairs."
    create_ec2_key_pair

    echo -e "\n[.] Step 5: Integrating scripts with your .bashrc"
    integrate_with_bashrc

    echo -e "\n[.] Step 6: Adding notification cron job"
    add_cron_job

    echo -e "\n[+] Ready to roll."
    echo
    echo "----------------------------------------------"
    echo
    echo -e "You can now try out our shiny new bash aliases:"
    echo -e "\t- check$EC2_NAME - To check an instance status"
    echo -e "\t- get$EC2_NAME - To get the instance IPv4 address."
    echo -e "\t- start$EC2_NAME - To start the instance."
    echo -e "\t- stop$EC2_NAME - To stop the instance."
    echo -e "\t- terminate$EC2_NAME - To terminate the instance."
    echo -e "\t- ssh$EC2_NAME - To ssh into that instance"
    echo
    echo "Go ahead, check them out. Start with creating EC2 first using 'start$EC2_NAME'"
    bash
fi
