#!/bin/bash

# title       :ct-aws-bandwith-limit
# description :This script permit to limit outgoing bandwidth on AWS service
# author      :Adrien DELLE CAVE (decryptus)
# date        :2018-02-17
# version     :0.1

PATH=/usr/local/sbin:/usr/local/bin:/sbin:/bin:/usr/sbin:/usr/bin

CT_AWS_BL_AWK_BIN="${CT_AWS_BL_AWK_BIN:-awk}"
CT_AWS_BL_CURL_BIN="${CT_AWS_BL_CURL_BIN:-curl}"
CT_AWS_BL_CUT_BIN="${CT_AWS_BL_CUT_BIN:-cut}"
CT_AWS_BL_GREP_BIN="${CT_AWS_BL_GREP_BIN:-grep}"
CT_AWS_BL_IP_BIN="${CT_AWS_BL_IP_BIN:-ip}"
CT_AWS_BL_IPT_BIN="${CT_AWS_BL_IPT_BIN:-iptables}"
CT_AWS_BL_JQ_BIN="${CT_AWS_BL_JQ_BIN:-jq}"
CT_AWS_BL_SED_BIN="${CT_AWS_BL_SED_BIN:-sed}"
CT_AWS_BL_SORT_BIN="${CT_AWS_BL_SORT_BIN:-sort}"
CT_AWS_BL_TC_BIN="${CT_AWS_BL_TC_BIN:-tc}"
CT_AWS_BL_XARGS_BIN="${CT_AWS_BL_XARGS_BIN:-xargs}"

CT_AWS_BL_IPS_URL="https://ip-ranges.amazonaws.com/ip-ranges.json"
CT_AWS_BL_REGION="${CT_AWS_BL_REGION:-eu-west-3}"
CT_AWS_BL_SERVICE="${CT_AWS_BL_SERVICE:-S3}"

CT_AWS_BL_IPT_COMMENT="CT: AWS ${CT_AWS_BL_SERVICE} Bandwidth Limit"
CT_AWS_BL_IFACE="${CT_AWS_BL_IFACE:-$(${CT_AWS_BL_IP_BIN} r s|${CT_AWS_BL_AWK_BIN} '/^default via/ {print $5}'|${CT_AWS_BL_SED_BIN} '/^$/d')}"
CT_AWS_BL_CLASS_ID="5"
CT_AWS_BL_BITRATE="${CT_AWS_BL_BITRATE:-300mbit}"

${CT_AWS_BL_TC_BIN} qdisc del dev ${CT_AWS_BL_IFACE} root
${CT_AWS_BL_TC_BIN} qdisc replace dev ${CT_AWS_BL_IFACE} root handle 1: htb
${CT_AWS_BL_TC_BIN} class replace dev ${CT_AWS_BL_IFACE} parent 1: classid 1:${CT_AWS_BL_CLASS_ID} htb rate ${CT_AWS_BL_BITRATE} prio 0
${CT_AWS_BL_TC_BIN} filter replace dev ${CT_AWS_BL_IFACE} parent 1: prio 0 protocol ip handle ${CT_AWS_BL_CLASS_ID} fw flowid 1:${CT_AWS_BL_CLASS_ID}
${CT_AWS_BL_TC_BIN} qdisc replace dev ${CT_AWS_BL_IFACE} parent 1:${CT_AWS_BL_CLASS_ID} sfq

${CT_AWS_BL_IPT_BIN} -t mangle -L OUTPUT --line-numbers|\
  ${CT_AWS_BL_GREP_BIN} "${CT_AWS_BL_IPT_COMMENT}"|\
  ${CT_AWS_BL_CUT_BIN} -d' ' -f1|\
  ${CT_AWS_BL_SORT_BIN} -r|\
  ${CT_AWS_BL_XARGS_BIN} --no-run-if-empty -n 1 ${CT_AWS_BL_IPT_BIN} -t mangle -D OUTPUT

for CT_AWS_BL_IP in `${CT_AWS_BL_CURL_BIN} -s "${CT_AWS_BL_IPS_URL}"|\
  ${CT_AWS_BL_JQ_BIN} -r ".prefixes[]|\
      select(.region==\"${CT_AWS_BL_REGION}\")|\
      select(.service==\"${CT_AWS_BL_SERVICE}\") | .ip_prefix"`
do
  ${CT_AWS_BL_IPT_BIN} -A OUTPUT -t mangle -d "${CT_AWS_BL_IP}" -j MARK --set-mark ${CT_AWS_BL_CLASS_ID} -m comment --comment "${CT_AWS_BL_IPT_COMMENT}"
done
