#!/bin/bash

# Configuration
BLOCK_IP_THRESHOLD=1
BLOCK_PREFIX_THRESHOLD=3
BLOCK_ASN_THRESHOLD=3

# Automatically remove old rules after X seconds
UNBLOCK_TIME=$((60 * 60 * 24 * 30)) # Unblock after 30 days

# IPs that should never be locked (can also be a prefix)
LAST_RESORT_IPS='
127.0.0.1
'

FILTER='
cat /var/log/secure | awk '\''$5~/^ssh/ && $6~/^Failed/ && $7~/^keyboard-interactive/ && $9~/[^invalid]/ {print $11}'\''
cat /var/log/secure | awk '\''$5~/^ssh/ && $6~/^Failed/ && $7~/^keyboard-interactive/ && $9~/[invalid]/ {print $13}'\''
'

# Path to the client script at the University of Bonn
BART_CLIENT='./request.py'

##########################################################
##### Nothing has to be changed after this line #####
##########################################################

IPTABLES=`which iptables`

# Test to see whether iptables supports comments; otherwise, no auto-unblock
UNBLOCK_AUTO=0

OUT=`cat /proc/net/ip_tables_matches | grep 'comment'`
if [ "$OUT" = "comment" ]; then
  UNBLOCK_AUTO=1
fi

function chain_exists
{
  TABLENAME="-t filter"
  CHAINNAME=$1
  if [ "$2" != "" ]; then
    TABLENAME="-t $1"
    CHAINNAME=$2
  fi
  if [ "$CHAINNAME" != "" ]; then
    ${IPTABLES} ${TABLENAME} -nL ${CHAINNAME} 2> /dev/null 1> /dev/null
    RET=$?
    return ${RET}
  fi
}

function delete_rule
{
  # Find correct rule parameters
  RULE=`echo ${@} | awk 'BEGIN{FS=" -j"}; {print $1}'`
  TARGET=`echo ${@} | awk 'BEGIN{FS=" -j"}; {print "-j"$2}'`

  ${IPTABLES} -S | grep --regexp="$RULE" | grep --regexp="$TARGET" | cut -d " " -f 2- | xargs -L1 ${IPTABLES} -D 2> /dev/null
}

function insert_rule
{
  # echo "insert: $@"
  delete_rule $@

  if [ "${UNBLOCK_AUTO}" -ne 1 ]; then
    ${IPTABLES} -I $@
  else
    TS=$(date +%s)
    ${IPTABLES} -I $@ -m comment --comment "Created: $TS"
  fi
}

function add_rule
{
  # echo "add: $@"
  delete_rule $@

  if [ "${UNBLOCK_AUTO}" -ne 1 ]; then
    ${IPTABLES} -A $@
  else
    TS=$(date +%s)
    ${IPTABLES} -A $@ -m comment --comment "Created: $TS"
  fi
}

if [ "$1" = "setup" ]; then
  # Deactivate comments to keep these rules permanently
  UNBLOCK_AUTO=0

  echo "Initialize iptables"
  # Create rule chains
  chain_exists LAST_RESORT || ${IPTABLES} -N LAST_RESORT
  ${IPTABLES} -F LAST_RESORT
  add_rule LAST_RESORT -j RETURN
  for IP in "${LAST_RESORT_IPS}"; do
    insert_rule LAST_RESORT -s ${IP} -j ACCEPT
  done

  chain_exists BLOCK_IP || ${IPTABLES} -N BLOCK_IP
  ${IPTABLES} -F BLOCK_IP
  add_rule BLOCK_IP -j RETURN

  chain_exists BLOCK_PREFIX || ${IPTABLES} -N BLOCK_PREFIX
  ${IPTABLES} -F BLOCK_PREFIX
  add_rule BLOCK_PREFIX -j RETURN

  chain_exists BLOCK_ASN || ${IPTABLES} -N BLOCK_ASN
  ${IPTABLES} -F BLOCK_ASN
  add_rule BLOCK_ASN -j RETURN


  # Create rules in the INPUT chain
  insert_rule INPUT -j BLOCK_IP
  insert_rule INPUT -j BLOCK_PREFIX
  insert_rule INPUT -j BLOCK_ASN
  insert_rule INPUT -j LAST_RESORT
  echo "Done"
  exit
elif [ "$1" = "teardown" ]; then
  echo "Remove all rules and chains"
  delete_rule INPUT -j BLOCK_IP
  delete_rule INPUT -j BLOCK_PREFIX
  delete_rule INPUT -j BLOCK_ASN
  delete_rule INPUT -j LAST_RESORT

  ${IPTABLES} -F LAST_RESORT 2> /dev/null
  ${IPTABLES} -X LAST_RESORT 2> /dev/null
  
  ${IPTABLES} -F BLOCK_IP 2> /dev/null
  ${IPTABLES} -X BLOCK_IP 2> /dev/null

  ${IPTABLES} -F BLOCK_PREFIX 2> /dev/null
  ${IPTABLES} -X BLOCK_PREFIX 2> /dev/null

  ${IPTABLES} -F BLOCK_ASN 2> /dev/null
  ${IPTABLES} -X BLOCK_ASN 2> /dev/null
  echo "Done"
  exit
fi

# Test whether setup is already performed
chain_exists BLOCK_IP || echo "Please start '$0 setup'"; exit

IPS=""
OLD_IFS=${iFS}
IFS="
"
for RULE in ${FILTER}; do
  # Apply the filter rule
  eval FOUND_IPS=\`${RULE}\`
  IPS="${IPS}
${FOUND_IPS}"
done

# Filter by IP address, incidence count, and use of the IP threshold
IPS=`echo "${IPS}" | grep -E '([0-9]{1,3}\.){3}[0-9]{1,3}' | sort -n | uniq -c | awk '$1 >= '${BLOCK_IP_THRESHOLD}' {print $2}'`

# echo "${IPS}"
for IP in ${IPS}; do
  insert_rule BLOCK_IP -s ${IP} -j DROP
done



# Finding and blocking the prefix
QUERY_RESPONSE=`echo "begin
prefix
$IPS
end" | netcat whois.cymru.com 43`

PREFIXES=`echo "${QUERY_RESPONSE}" | awk 'BEGIN{FS="|"}; {print $1" "$3}' | grep -v '^$' | sort -n | uniq -c | awk '$1 >= '${BLOCK_PREFIX_THRESHOLD}' {print $2" "$3}'`

# echo "${PREFIXES}"
for P in $( echo "${PREFIXES}" | awk '{print $2}' ); do
  DEL_IPS=`echo "${QUERY_RESPONSE}" | grep "${P}" | awk 'BEGIN{FS="|"}; {print $2}' | sort -n | uniq | awk '{print $1}'`
  for IP in ${DEL_IPS}; do
    # echo "Delete single check for IP: ${IP}"
    delete_rule BLOCK_IP -s ${IP} -j DROP
  done
  insert_rule BLOCK_PREFIX -s ${P} -j DROP
done



# Filter all prefixes of an AS
ASNLIST=`echo "${PREFIXES}" | awk '{print $1}' | sort -n | uniq -c | awk '$1 >= '${BLOCK_ASN_THRESHOLD}' {print $2}'`

if [ -x "${BART_CLIENT}" ]; then
  for AS in ${ASNLIST}; do
    PREFIX_LIST=`${BART_CLIENT} announced_prefixes $AS`
    for P in ${PREFIX_LIST}; do
      delete_rule BLOCK_PREFIX -s ${P} -j DROP
      insert_rule BLOCK_ASN -s ${P} -j DROP
    done
  done
fi



# Discard old entries
if [ "${UNBLOCK_AUTO}" -eq 1 ]; then
  NOW=$(date +%s)
  RULES=`${IPTABLES} -S | cut -d " " -f 2-`
  for R in $RULES; do
    TIMESTAMP=`echo "${R}" | grep -o '"Created: .*"' | tr -d '"' | awk '{print $2}'`
    if [ "$TIMESTAMP" != "" ] && [ "$TIMESTAMP" -lt $(($NOW - $UNBLOCK_TIME)) ]; then
      delete_rule $R
    fi
  done
fi
IFS=${OLD_IFS}
