#!/usr/bin/env bash
# Author: Zhang Huangbin <zhb@iredmail.org>
# Puprpose:
#   - Store banned IP address in SQL db while it's banned by Fail2ban.
#   - Remove unbanned IP address from SQL db while it's unbanned by Fail2ban.
#   - Unban IP addresses which have column `remove=1`.

# Usage:
#
#   *) Store a banned IP address:
#
#       banned_db ban <ip> <ports> <protocol> <jail> <failures> [loglines]
#
#       - <ip>: One IP address each time.
#       - <ports>: Network ports. Multiple ports must be separated by comma.
#       - <protocol>: `tcp` or `udp`.
#       - <jail>: Fail2ban jail name.
#       - <failures>: number of times the failure occurred in the log file.
#       - [loglines]: matched log lines. OPTIONAL.
#
#   *) Remove an one or multiple unbanned IP addresses. Notes:
#
#       - it removes IP from all jails.
#       - multiple IP addresses must be separated by space.
#
#       banned_db unban <ip> [ip] [ip]
#
#   *) Cleanup a jail. When Fail2ban is stopping or restarting, `cleanup` will
#      be executed. Cleanup manually is supported too:
#
#       banned_db cleanup <jail>
#
#   *) Query SQL db and remove IP addresses which have `remove=1`.
#
#       banned_db unban_db
#
# Examples:
#
#   banned_db start dovecot-iredmail
#   banned_db ban   192.168.0.1 110,143,993,995 tcp dovecot-iredmail 3
#   banned_db unban 192.168.0.1
#   banned_db stop dovecot-iredmail
#   banned_db unban_db

#
# Sample Fail2ban jail config file (/etc/fail2ban/jail.d/xx.local):
#
#   [jail-name]
#   ...
#   action = ...[your other actions here]...
#            banned_db[name=jail-name, port="80", protocol=tcp]
#
# WARNING: the name set in `banned_db[name=]` must be same as the jail name.

export PATH="/usr/bin:/usr/local/bin:$PATH"

export DB_NAME="fail2ban"
export DB_TABLE_BANNED="banned"
export DB_TABLE_JAILS="jails"
export DB_USER="fail2ban"

# GeoIP
export CMD_GEOIPLOOKUP="$(which geoiplookup 2>/dev/null)"
export CMD_GEOIPLOOKUP6="$(which geoiplookup6 2>/dev/null)"

# `dig`. Used to query reverse hostname.
export CMD_DIG="/usr/bin/dig"

if [ -f /root/.my.cnf-fail2ban ]; then
    export CMD_SQL="mysql --defaults-file=/root/.my.cnf-fail2ban ${DB_NAME}"
    export DB_TYPE="mysql"
elif [ -f /root/.my.cnf ]; then
    export CMD_SQL="mysql --defaults-file=/root/.my.cnf ${DB_NAME}"
    export DB_TYPE="mysql"
else
    # Absolute path to ~/.pgpass
    #   - RHEL:             /var/lib/pgsql/.pgpass
    #   - Debian/Ubuntu:    /var/lib/postgresql/.pgpass
    #   - FreeBSD:          /var/db/postgres/.pgpass
    #   - OpenBSD:          /var/postgresql/.pgpass
    for dir in \
        /var/lib/pgsql \
        /var/lib/postgresql \
        /var/db/postgres \
        /var/postgresql; do
        if [ -f ${dir}/.pgpass ]; then
            export PGPASSFILE="${dir}/.pgpass"
            export CMD_SQL="psql -U ${DB_USER} -d ${DB_NAME}"
            export DB_TYPE="pgsql"
            break
        fi
    done
fi

if [ X"${CMD_SQL}" == X'' ]; then
    echo "No MySQL or PostgreSQL related config file found. Abort."
    echo " - MySQL: /root/.my.cnf-fail2ban (or /root/.my.cnf)"
    echo " - PostgreSQL: ~/.pgpass (under PostgreSQL daemon user's home directory)"
    exit 255
fi

export _action="$1"

if [[ X"${_action}" == X'start' ]]; then
    _jail="${2}"

    if [[ X"${DB_TYPE}" == X'pgsql' ]]; then
        # CentOS 7 ships PostgreSQL-9.2 which doesn't support `ON CONFLICT DO NOTHING`,
        # so we query it first, insert it if not existing.
        (${CMD_SQL} <<EOF
SELECT id FROM ${DB_TABLE_JAILS} WHERE name='${_jail}' LIMIT 1;
EOF
) | grep '1 row' &>/dev/null

        if [[ X"$?" != X'0' ]]; then
            ${CMD_SQL} >/dev/null <<EOF
INSERT INTO ${DB_TABLE_JAILS} (name, enabled) VALUES ('${_jail}', 1);
EOF
        fi
    else
        ${CMD_SQL} >/dev/null <<EOF
INSERT IGNORE INTO ${DB_TABLE_JAILS} (name, enabled) VALUES ('${_jail}', 1);
EOF
    fi
elif [[ X"${_action}" == X"ban" ]]; then
    _ip="${2}"
    _ports="${3}"
    _protocol="${4}"
    _jail="${5}"
    _failures="${6}"

    shift 6

    # (base64) Encode log lines to avoid possible SQL injection.
    if [[ -x /usr/bin/base64 ]]; then
        # Linux
        _loglines="$(echo $@ | base64)"
    elif [[ -x /usr/bin/b64encode ]]; then
        # OpenBSD and FreeBSD.
        _loglines="$(echo $@ | b64encode - |grep -Ev '^(begin-base64 |====$)')"
    else
        _loglines=""
    fi

    if [[ X"${_ip}" == X'' ]] || \
        [[ X"${_ports}" == X'' ]] || \
        [[ X"${_protocol}" == X'' ]] || \
        [[ X"${_jail}" == X'' ]]; then
        echo "IP, ports, protocol, or jail name is empty. Abort."
        exit 255
    fi

    _hostname="$(hostname)"

    # Lookup for country name.
    _country=''
    if echo ${_ip} | grep ':' &>/dev/null; then
        if [[ -x ${CMD_GEOIPLOOKUP6} ]]; then
            _country="$(${CMD_GEOIPLOOKUP6} ${_ip} | grep '^GeoIP Country Edition:' | awk -F': ' '{print $2}' | grep -iv 'not found' | tr -d '"' | tr -d "'")"
        fi
    else
        if [[ -x ${CMD_GEOIPLOOKUP} ]]; then
            _country="$(${CMD_GEOIPLOOKUP} ${_ip} | grep '^GeoIP Country Edition:' | awk -F': ' '{print $2}' | grep -iv 'not found' | tr -d '"' | tr -d "'")"
        fi
    fi

    # Lookup reverse DNS name.
    _rdns=''
    if [[ -x "${CMD_DIG}" ]]; then
        _rdns_orig="$(${CMD_DIG} +short +timeout=3 -x ${_ip} 2>/dev/null)"
        _rdns_strip="${_rdns_orig%\.}"
        printf -v _rdns "${_rdns_strip}"
    fi

    if [ X"${DB_TYPE}" == X'mysql' ]; then
        # MySQL stores local time with time zone info, we expect UTC time.
        ${CMD_SQL} >/dev/null <<EOF
INSERT IGNORE INTO ${DB_TABLE_BANNED} (
        ip, ports, protocol, jail,
        hostname, country, rdns, timestamp,
        failures, loglines)
    VALUES (
        '${_ip}', '${_ports}', '${_protocol}', '${_jail}',
        '${_hostname}', '${_country}', '${_rdns}', UTC_TIMESTAMP(),
        '${_failures}', '${_loglines}');
EOF
    else
        # CentOS 7 ships PostgreSQL-9.2 which doesn't support `ON CONFLICT DO NOTHING`,
        # so we query it first, insert it if not existing.
        (${CMD_SQL} <<EOF
SELECT id FROM ${DB_TABLE_BANNED} WHERE ip='${_ip}' AND ports='${_ports}' AND protocol='${_protocol}' AND jail='${_jail}' LIMIT 1;
EOF
) | grep '1 row' &>/dev/null

        if [[ X"$?" == X'0' ]]; then
            echo "Already banned."
        else
            ${CMD_SQL} >/dev/null <<EOF
INSERT INTO ${DB_TABLE_BANNED} (
        ip, ports, protocol, jail,
        hostname, country, rdns,
        failures, loglines)
    VALUES (
        '${_ip}', '${_ports}', '${_protocol}', '${_jail}',
        '${_hostname}', '${_country}', '${_rdns}',
        '${_failures}', '${_loglines}');
EOF
            echo "Stored."
        fi
    fi

elif [[ X"${_action}" == X"unban" ]]; then
    shift 1
    _ips="$@"

    if [[ X"${_ips}" == X'' ]]; then
        echo "No IP address(es) specified."
    else
        for _ip in ${_ips}; do
            ${CMD_SQL} >/dev/null <<EOF
DELETE FROM ${DB_TABLE_BANNED} WHERE ip='${_ip}';
EOF
        done

        [[ X"$?" == X'0' ]] && echo "Removed."
    fi

elif [[ X"${_action}" == X'stop' ]] || [[ X"${_action}" == X"cleanup" ]]; then
    _jail="$2"

    if [[ X"${_jail}" != X'' ]]; then
        ${CMD_SQL} >/dev/null <<EOF
DELETE FROM ${DB_TABLE_BANNED} WHERE jail='${_jail}';
DELETE FROM ${DB_TABLE_JAILS} WHERE name='${_jail}';
EOF
    fi

    [[ X"$?" == X'0' ]] && echo "All IP addresses have been removed."
elif [[ X"${_action}" == X"unban_db" ]]; then
    # Call fail2ban-client to unban given IP address(es).
    tmp_file="$(mktemp)"

    # Exclude extra info on output, just leave jail/ip.
    (${CMD_SQL} <<EOF
SELECT jail, ip FROM ${DB_TABLE_BANNED} WHERE remove=1;
EOF
) | grep -Ev '(jail.*ip|\--|\(|^$)' | tr -d '|' >> ${tmp_file}

    while read jail ip; do
        # Avoid SQL injection: don't allow whitespace, ';', quotes in
        # jail name and IP address.
        if echo ${jail} | grep "[ ;\"\']" &>/dev/null; then
            echo "[WARNING] Invalid jail name: '${jail}'."
            continue
        fi

        if echo ${ip} | grep "[ ;\"\']" &>/dev/null; then
            echo "[WARNING] Invalid IP address: '${ip}'."
            continue
        fi

        # fail2ban-client returns number of processed rows on command line,
        # let's discard it to avoid noise/confusion.
        fail2ban-client set ${jail} unbanip ${ip} >/dev/null
        [[ X"$?" == X'0' ]] && echo "Unbanned ${ip} from jail [${jail}]."
    done < ${tmp_file}
    rm -f ${tmp_file} &>/dev/null
fi