A simple script to protect system from DOS attack

autoban.sh

#!/bin/bash
# A simple script to protect system from DOS attack
# create a crontab conf file /etc/cron.d/autoban like this:
# * * * * *       root /sbin/autoban.sh ban >/dev/null 2>&1 &
# */10 * * * *    root /sbin/autoban.sh unban - 5 >/dev/null 2>&1 &
# 01 * * * *      root /sbin/autoban.sh unban 6 20 >/dev/null 2>&1 &
# 01 09 * * *     root /sbin/autoban.sh unban 21 - >/dev/null 2>&1 &
#
PATH=/bin:/usr/bin:/usr/sbin:/sbin
export PATH

[ -f /etc/autoban.conf ] && . /etc/autoban.conf
MAXCONN=${MAXCONN-150}
PORT=${PORT-80}

BLACKLIST=/etc/autoban_black_list
WHITELIST=/etc/autoban_white_list

iptables=$(which iptables)
ss=$(which ss)

_add() {
    local i=$1
    local f=$2
    if grep -wq "^$i .*$" "$f"; then
        local n=$(awk "/^$i / {print \$2}" $f)
        local nn=$((n+1))
        sed -i "s/^$i .*/$i $nn/" $f
    else
        echo "$i 1" >>"$f"
    fi
}

_in() {
    local i=$1
    local f=$2
    if grep -wq "^$i" "$f"; then
        return 0
    fi
    return 1
}

getips() {
    $ss -an "( sport = :$PORT )" | \
        awk '{print $NF}' | \
        sed 's/^::ffff://' | \
        awk --posix -F':' '/^[0-9]{1,}\./ {print $1}' | \
        sort | uniq -c
}

ban() {
    local iptmp=$(mktemp -p /dev/shm)
    local wltmp=$(mktemp -p /dev/shm)

    getips >$iptmp
    [ -f $BLACKLIST ] || touch $BLACKLIST
    [ -f $WHITELIST ] || touch $WHITELIST
    grep -v '^#' $WHITELIST >$wltmp
    local wlip
    for wlip in $(awk '{print $1}' $wltmp); do
        sed -i "/$wlip/d" $iptmp
    done
    sed -i '/127.0.0.1/d' $iptmp

    local conn ip
    while read conn ip; do
        if [ $conn -ge $MAXCONN ]; then
            if _in $ip $BLACKLIST; then
                :
            else
                $iptables -I INPUT -s $ip -p tcp --dport $PORT -j DROP >/dev/null 2>&1
            fi
            _add $ip $BLACKLIST
        fi
    done <$iptmp
    rm -f $iptmp $wltmp
}

unban() {
    local rank_start=$1
    local rank_end=$2
    local rstart=0
    local rend=9999
    if is_digit $rank_start; then
        rstart=$rank_start
    fi
    if is_digit $rank_end; then
        rend=$rank_end
    fi
    local ip rank
    if ! [ -f $BLACKLIST ]; then
        return 0
    fi
    while read ip rank; do
        if [[ $rank -ge $rstart ]] && [[ $rank -le $rend ]]; then
            $iptables -D INPUT -s $ip -j DROP &>/dev/null 2>&1
            sed -i "/^$ip $rank$/d" $BLACKLIST
            sleep .1
        fi
    done <$BLACKLIST
}

usage() {
    echo "usage: autoban <action> [options]
options:
    ban         : ban ip
    unban r1 r2 : unban ip between rank1 to rank2
    list        : output client ip list
"
}

if [ $# -lt 1 ]; then
    usage
    exit 2
fi
act=$1
case $act in
    'unban')
        if [ $# -ne 3 ]; then
            usage
            exit 2
        fi
        unban $2 $3 >/dev/null 2>&1
        exit
        ;;
    'ban')
        ban >/dev/null 2>&1
        ;;
    'list')
        ips=$(getips)
        echo total: $(echo "$ips" | wc -l)
        echo top 10:
        echo "$ips" | sort -n | tail -10
        ;;
    *)
        usage
        exit
        ;;
esac