#!/bin/bash -e
#
# Amazon AWS blocker through iptables.
#
# First we use curl to grab the official list of ranges from Amazon. The -s
# prevents extraneous output from curl, and the -L makes it follow redirects.
#
# The ranges are passed to jq, a JSON parser. The -r makes jq output raw data
# without quotes. We only need the list of prefixes, so we discard everything
# else.

POSITION=1
FILTERS=""
JSON_URL="https://ip-ranges.amazonaws.com/ip-ranges.json"


# Get the line where the jump will be inserted at.
# Useful if you want e.g related / established rules for outgoing traffic.
if [[ -n $1 ]]; then
    POSITION=$1
    shift
fi


##
# Builds region filters based on CLI arguments
#
# Arguments: CLI arguments as passed by $*
#
function build_filters() {
    for arg in ${@:1}; do
        if [[ -n $filters ]]; then
            filters=$filters", "
        fi

        filters=$filters"select(.region | contains(\"$arg\"))"
    done

    if [[ -n $filters ]]; then
        filters=" | "$filters
    fi

    echo $filters
}


##
# Extracts IP ranges from an Amazon JSON file
#
# Arguments:
#     $1 AWS JSON content
#     $2 Prepared filter string
#     $3 Group to extract IP ranges from (e.g. prefixes)
#     $4 Object key for IP ranges (e.g ip_prefix)
#
function extract_ip_ranges() {
    local json=$1
    local filters=$2
    local array=$3
    local prefix=$4

    local group='group_by(.'$prefix')'
    local map='map({ "ip": .[0].'$prefix', "regions": map(.region) | unique, "services": map(.service) | unique })'

    local to_string='.ip + " \"" + (.regions | sort | join (", ")) + "\" \"" + (.services | sort | join (", ")) + "\""'
    local process='[ .'$array"[]$filters ] | $group | $map | .[] | $to_string"

    local ranges=$(echo "$json" | jq -r "$process" | sort -Vu)
    echo "$ranges"
}


##
# Creates the AWS iptables chain if it doesn't exist, then flushes it
#
# Arguments:
#     $1 Version to use. Omit for v4
#     $2 Position to insert chain statement at
#
function create_and_flush_chain() {
    local version=$1
    local position=$2
    local cmd=ip${version}tables

    $cmd -n --list AWS >/dev/null 2>&1 \
        || ($cmd -N AWS && $cmd -I INPUT $position -j AWS)

    $cmd -F AWS
}


##
# Adds an iptables rule for each line in ranges
#
# Arguments:
#     $1 Version to use. Omit for v4
#     $2 Prepared lines
#
function add_iptables_rules() {
    local version=$1
    local cmd=ip${version}tables
    local lines
    local data

    IFS=$'\n' lines=($2)
    unset IFS

    for line in "${lines[@]}"; do
        eval local data=($line)
        local ip=${data[0]}
        local regions=$(echo ${data[1]} | tr '[:upper:]' '[:lower:]')
        local services=$(echo ${data[2]} | tr '[:upper:]' '[:lower:]')

        $cmd -A AWS -s "$ip" -j REJECT -m comment --comment "$regions = $services"
    done
}



# Retrieve IP ranges definition
# Either from an URL or file input (e.g. "< ranges.json")
if [ ! -t 0 ]; then
    JSON=$(cat - <&0)
else
    JSON=$(curl -s -L $JSON_URL)
fi

FILTERS=$(build_filters "$*")


# IPv4
create_and_flush_chain "" $position
V4_RANGES=$(extract_ip_ranges "$JSON" "$FILTERS" "prefixes" "ip_prefix")
add_iptables_rules ""  "$V4_RANGES"


# IPv6
create_and_flush_chain 6 $position
V6_RANGES=$(extract_ip_ranges "$JSON" "$FILTERS" "ipv6_prefixes" "ipv6_prefix")
add_iptables_rules "6" "$V6_RANGES"