#!/bin/bash

#
# Adds Cloudflare IPV4 addresses to the Inbound rules of a AWS Security Group
#

## defining helper functions first

# show usage
show_usage() {
    echo "usage: $0 -r <REGION> -s <SECURITY-GROUP-ID> -p <tcp|udp|icmp> -t <PORT-NUMBER> -n <SNS-TOPIC>"
}

# log message to STDOUT
log() {
    echo "[$(date)]: $*"
}

# wraps command execution to log possible failures and publish sns topics
execute_command() {
    CMD="$1"
    TMP_FILE="$2"
    SNS_TOPIC="$3"
    SEND_ON_SUCCESS="$4"
    log "issuing $CMD"
    OUTPUT=$($CMD 2>&1 > $TMP_FILE)
    CMD_EXIT_STATUS=$?
    if [ $CMD_EXIT_STATUS -ne 0 ]; then
        MESSAGE="$CMD failed with status $CMD_EXIT_STATUS: $OUTPUT"
        log $MESSAGE 
        [ $SNS_TOPIC != false ] && aws sns publish --topic-arn $SNS_TOPIC --message "$MESSAGE" --subject "[INFRA] - $0 failed"
        rm -fv $TMP $ALLOWED_IP_RANGES $CLOUDFLARE_IP_RANGES
        exit 1
    elif [ $SEND_ON_SUCCESS != false ]; then
       MESSAGE="$CMD executed sucessfully"
       [ $SNS_TOPIC != false ] && aws sns publish --topic-arn $SNS_TOPIC --message "$MESSAGE" --subject "[INFRA] - $0 successful output"
    fi 
}


## MAIN
# ARGS parse
while getopts ":r:s:p:t:n:" opt; do
    case $opt in
        r)
            REGION=$OPTARG
            ;;
        s)
            SECURITY_GROUP_ID=$OPTARG
            ;;
        p)
            PROTOCOL=$OPTARG
            ;;
        t)
            PORT=$OPTARG
            ;;
        n)
            SNS_TOPIC=$OPTARG
            ;;
        ?)
            show_usage
            exit 1
            ;;
        :)
            echo "Option -$OPTARG requires an argument." >&2
            show_usage
            exit 1
            ;;
  esac
done

# test for required args
if [ -z $SECURITY_GROUP_ID ] || [ -z $PROTOCOL ] || [ -z $PORT ]; then
    show_usage
    exit 1
fi

# if SNS_TOPIC is not passed, set to false
if [ -z $SNS_TOPIC ]; then
    SNS_TOPIC=false
fi

# if REGION is not passed, set to us-east-1
if [ -z $REGION ]; then
    REGION="us-east-1"
fi

TMP=$(mktemp)
ALLOWED_IP_RANGES=$(mktemp)
CLOUDFLARE_IP_RANGES=$(mktemp)

# Get current ingress rules for port 80, protocol tcp for the given security group id (sorted by ip address range)
execute_command "aws ec2 describe-security-groups --group-id $SECURITY_GROUP_ID --region $REGION" $TMP $SNS_TOPIC false

cat $TMP | jq ".SecurityGroups[].IpPermissions[] | select(.FromPort == ${PORT} and .ToPort == ${PORT} and .IpProtocol == \"${PROTOCOL}\") | .IpRanges[] | .CidrIp" | tr -d '"' | sort -t . -k 1,1n -k 2,2n -k 3,3n -k 4,4n > $ALLOWED_IP_RANGES 
log "security group ip ranges allowed for protocol $PROTOCOL and port $PORT: $(cat $ALLOWED_IP_RANGES)"

# Get Cloudflare IP Ranges (sorted by ip address range)
execute_command "curl -sS https://www.cloudflare.com/ips-v4" $TMP $SNS_TOPIC false

cat $TMP | sort -t . -k 1,1n -k 2,2n -k 3,3n -k 4,4n > $CLOUDFLARE_IP_RANGES
log "cloudflare ip ranges: $(cat $CLOUDFLARE_IP_RANGES)"

# Get the missing IP ranges and add them to the security group
log "issuing diff -u between security group existing ip addresses and cloudflare addresses: diff -u $ALLOWED_IP_RANGES $CLOUDFLARE_IP_RANGES"
# Add missing cloudflare ip ranges to the security group allowed ingress rule table
MISSING_CF_IP_RANGES=$(diff -u $ALLOWED_IP_RANGES $CLOUDFLARE_IP_RANGES | grep '^+' | grep -v '+++' | tr -d '+')
log "missing Cloudflare IP Ranges: $MISSING_CF_IP_RANGES"
for IP_RANGE in $MISSING_CF_IP_RANGES; do
    execute_command "aws ec2 authorize-security-group-ingress --group-id $SECURITY_GROUP_ID --protocol $PROTOCOL --port $PORT --cidr $IP_RANGE --region $REGION" $TMP $SNS_TOPIC true
done

log "clean up"
rm -fv $TMP $ALLOWED_IP_RANGES $CLOUDFLARE_IP_RANGES