#!/bin/bash

# Author:  Zilvinas Verseckas
# Company: Kolmisoft
# Year:    2016
# About:   Blocked Countries manager
#
#          Script can be run in two ways with the 'UPDATE' or 'RENEW' parameter.
#
#          RENEW CASE:
#          All zone files in /usr/local/m2/ip_zones will be deleted, ipset will be restarted, iptables
#          chains will be flushed and deleted. This way the next update will run the download and create
#          everyting anew.
#
#          UPDATE CASE:
#          Retrieves the ISO 3166 codes list for Countries. The list is then compared
#          with the existing ipsets and split into two lists: the ADD_LIST (ipsets to be added or updated)
#          and the DESTROY_LIST (ipsets to be droppped).
#
#          For each of the DESTROY_LIST ipset names (m2_{code}_ips) kernel relations in iptables and the
#          ipsets themselves will be destroyed.
#
#          For each of the ADD_LIST ipset names the {code}.zone file from ideny.com will be downloaded
#          (if not present) and the corresponding ipset will be populated with the file's content.
#
#          Each time the iptables will be checked to have a default whitelist which includes a Kolmisoft
#          ip, an ipdeny.com ip, DNS, local and server ips from the database (for misconfiguration cases).
#
#          Blacklist rejects all protocols and ports from the blocked countries except for SSH p22
#

# v 1.2.1

. /usr/src/m2/framework/bash_functions.sh
k_iptables_locking_option

if [[ "`/sbin/pidof -x $(basename $0) -o %PPID`" ]]; then
    report "$(basename $0) script is already running with PID `/sbin/pidof -x $(basename $0) -o %PPID`" 1
    exit
fi

# Database credentials
DB_PASSWORD=`cat /etc/m2/system.conf | grep dbsecret | awk '{print $3}'`
DB_USERNAME=`cat /etc/m2/system.conf | grep dbuser | awk '{print $3}'`
DB_NAME=`cat /etc/m2/system.conf | grep dbname | awk '{print $3}'`
DB_HOST=`cat /etc/m2/system.conf | grep dbhost | awk '{print $3}'`

# Chain names
BLACKLIST_CHAIN="M2-BLOCKED-COUNTRIES"

# List of ipsets to update/destroy
ADD_LIST=""
DESTROY_LIST=""

format_date() {
    date "+%Y-%m-%d %H:%M:%S"
}

retrieve_data() {
    # Get marked country codes in ISO 3166 standard
    local codes=`/usr/bin/mysql -h $DB_HOST $DB_NAME -u $DB_USERNAME -p$DB_PASSWORD -s -N -e "SELECT iso3166code FROM blocked_countries JOIN directions ON direction_id = directions.id WHERE state = 1 AND iso3166code != '00'" | tr "[A-Z]" "[a-z]"`

    # Handle MySQL errors
    if [[ "$codes" == *"error"* ]]; then
        report "Unexpected MySQL error: ${codes}! Exiting" 6
        exit 1
    fi

    # Form a list of ipset names: "m2_${code}_ips"
    local ip_list_names=""
    for code in $codes; do
        ip_list_names+=" m2_${code}_ips"
    done

    # Synchronize the database list and the ipset list
    local unique=$(echo "$ip_list_names" `/usr/sbin/ipset -L -n` | tr " " "\n" | sort | uniq)
    for name in $unique; do
        if [[ "$ip_list_names" == *"$name"* ]]; then
            ADD_LIST+=" ${name}"
        elif [[ "$name" == "m2_"* ]]; then
            DESTROY_LIST+=" ${name}"
        fi
    done

    report "IP sets to add: [${ADD_LIST} ]" 3
    report "IP sets to destroy: [${DESTROY_LIST} ]" 3
}

# Clean up the iptables and the ipsets
clean_old_relations() {
    if [ "$DESTROY_LIST" == "" ]; then
        return
    fi

    if /sbin/iptables -n -L $BLACKLIST_CHAIN >/dev/null 2>&1; then
        # Delete all rules containing the ipset name
        local regexp=`echo "$DESTROY_LIST" | sed -e 's/^[ ]*//' | tr " " "|"`
        local rules=`/sbin/iptables $l_opt -n -L $BLACKLIST_CHAIN --line-numbers | grep -E "$regexp" | awk '{ print $1 }' | sort -r -g | tr " " "\n"`

        for rule in $rules; do
            /sbin/iptables $l_opt -D $BLACKLIST_CHAIN $rule 2>&1
        done
    fi

    for ip_set in $DESTROY_LIST; do
        report "Deleting IP set: ${ip_set}" 2
        /usr/sbin/ipset destroy $ip_set -exist
    done
}

# Create a whitelist for solving possible misconfigurations
add_default_whitelist() {
    # Allow server ips from M2
    local server_ips=`/usr/bin/mysql -h $DB_HOST $DB_NAME -u $DB_USERNAME -p$DB_PASSWORD -s -N -e "SELECT server_ip FROM servers"`

    for ip in $server_ips; do
        if ! /sbin/iptables $l_opt -L $BLACKLIST_CHAIN -v -n | grep "M2 Server ${ip}" | grep -q ACCEPT; then
            /sbin/iptables $l_opt -I $BLACKLIST_CHAIN -s $ip -j ACCEPT -m comment --comment "M2 Server ${ip}"
        fi
    done

    # Allow the Server network interface ips
    for ip in `ip addr show | grep "inet " | tr -s " " | cut -d " " -f3`; do
        if [[ "$ip" == *"127.0.0"* ]]; then
            continue
        fi

        if ! /sbin/iptables $l_opt -L $BLACKLIST_CHAIN -v -n | grep  "Local ${ip}" | grep -q ACCEPT; then
            /sbin/iptables $l_opt -I $BLACKLIST_CHAIN -s $ip -j ACCEPT -m comment --comment "Local ${ip}"
        fi
    done

    # Allow the DNS servers
    for ip in `cat /etc/resolv.conf | grep -E "nameserver [0-9]+\.[0-9]+\.[0-9]+\.[0-9]+" | cut -f2 -d " "`; do
        if ! /sbin/iptables $l_opt -L $BLACKLIST_CHAIN -v -n | grep "DNS Server ${ip}" | grep -q ACCEPT; then
            /sbin/iptables $l_opt -I $BLACKLIST_CHAIN -s $ip -j ACCEPT -m comment --comment "DNS Server ${ip}"
        fi
    done

    # Allow ipdeny API
    if ! /sbin/iptables $l_opt -L $BLACKLIST_CHAIN -v -n | grep "API for ipdeny.com" | grep -q ACCEPT; then
        /sbin/iptables $l_opt -I $BLACKLIST_CHAIN -s 192.241.240.22 -j ACCEPT -m comment --comment "API for ipdeny.com"
    fi

    # Allow kolmisoft network connections
    if ! /sbin/iptables $l_opt -L $BLACKLIST_CHAIN -v -n | grep "For Kolmisoft access" | grep -q ACCEPT; then
        /sbin/iptables $l_opt -I $BLACKLIST_CHAIN -s 77.240.248.90 -j ACCEPT -m comment --comment "For Kolmisoft access"
    fi

    if ! /sbin/iptables $l_opt -L $BLACKLIST_CHAIN -v -n | grep "For Kolmisoft VPN access" | grep -q ACCEPT; then
        /sbin/iptables $l_opt -I $BLACKLIST_CHAIN -s 157.90.244.37 -j ACCEPT -m comment --comment "For Kolmisoft VPN access"
    fi

    if ! /sbin/iptables $l_opt -L $BLACKLIST_CHAIN -v -n | grep -F "For kolmisoft.com access" | grep -q ACCEPT; then
        /sbin/iptables $l_opt -I $BLACKLIST_CHAIN -s 5.9.57.125 -j ACCEPT -m comment --comment "For kolmisoft.com access"
    fi

    if ! /sbin/iptables $l_opt -L $BLACKLIST_CHAIN -v -n | grep -F "For kolmisoft_c access" | grep -q ACCEPT; then
        /sbin/iptables $l_opt -I $BLACKLIST_CHAIN -s 116.203.27.119 -j ACCEPT -m comment --comment "For kolmisoft_c access"
    fi

    if ! /sbin/iptables $l_opt -L $BLACKLIST_CHAIN -v -n | grep -F "For monitoring.kolmisoft.com access" | grep -q ACCEPT; then
        /sbin/iptables $l_opt -I $BLACKLIST_CHAIN -s 135.181.92.61 -j ACCEPT -m comment --comment "For monitoring.kolmisoft.com access"
    fi

    # Allow all SSH connections
    if ! /sbin/iptables $l_opt -L $BLACKLIST_CHAIN -v -n | grep "Allow SSH" | grep -q ACCEPT; then
        /sbin/iptables $l_opt -I $BLACKLIST_CHAIN -p tcp --dport 22 -m conntrack --ctstate NEW,ESTABLISHED -j ACCEPT -m comment --comment "Allow SSH"
    fi
}

remove_chain() {
    local chain=$1
    if /sbin/iptables $l_opt -n -L $chain >/dev/null 2>&1; then
        /sbin/iptables $l_opt -D INPUT -j $chain >/dev/null 2>&1
        /sbin/iptables $l_opt -F $chain
        /sbin/iptables $l_opt -X $chain
        report "Chain '${chain}' removed" 3
    fi
}

add_chain() {
    # If a blacklist chain is not yet added, add it
    if ! /sbin/iptables $l_opt -n -L $BLACKLIST_CHAIN >/dev/null 2>&1; then
        # Create it
        /sbin/iptables $l_opt -N $BLACKLIST_CHAIN
        report "Chain '${BLACKLIST_CHAIN}' created" 3
    fi
}

# Update the IP sets and add them to iptables
add_rule() {
    for ip_set in $ADD_LIST; do

        report "Adding IP set ${ip_set}" 3

        local code=`echo $ip_set | cut -f2 -d_`
        local path=/usr/local/m2/ip_zones/${code}.zone

        # Create a new IP set if not present and fill it with IPs from a zone file
        if /usr/sbin/ipset -L -n | grep -q $ip_set; then
            report "IP set ${ip_set} already exists" 3
        else
            if [ -f $path ]; then
                report "Found zone file for country with code: ${code} in ${path}." 3
            else
                report "Donwloading zone file for country with code: ${code} and saving it in ${path}" 3
                wget "http://www.ipdeny.com/ipblocks/data/aggregated/${code}-aggregated.zone" -O "$path" -q --timeout=5 -t 1
            fi

            if [[ ! -s $path ]]; then
                report "Download for zone file ${code}.zone failed" 1
                rm -f $path >/dev/null
                echo ""
                continue
            else
                /usr/sbin/ipset create $ip_set hash:net -exist

                report "Inserting IPs into IP set ${ip_set}" 3

                for ip in `cat $path`; do
                    /usr/sbin/ipset add $ip_set $ip -exist
                done
            fi
        fi

        # Add rules to a corresponding chain if do not exist
        if /sbin/iptables $l_opt -n -L $BLACKLIST_CHAIN | grep -q $ip_set; then
            report "Iptables relation for IP set ${ip_set} already exists" 3
            echo ""
            continue
        fi

        report "Inserting IP set ${ip_set} into iptables! Country with code: ${code} will be blocked!" 3
        /sbin/iptables $l_opt -A $BLACKLIST_CHAIN -m set --set $ip_set src -j DROP
        echo ""
    done

    if [[ ! `/usr/sbin/ipset -L -n` ]]; then
        remove_chain $BLACKLIST_CHAIN
        report "No ipsets created. Exiting" 3
        exit 0
    fi

    # Append the chain to INPUT
    if ! /sbin/iptables $l_opt -n -L INPUT | grep -q $BLACKLIST_CHAIN; then
        /sbin/iptables $l_opt -A INPUT -j $BLACKLIST_CHAIN
        report "Chain '${BLACKLIST_CHAIN}' added to INPUT chain" 3
    fi
}

if [ "$1" == "UPDATE" ]; then

    report "[$(format_date)] Updating Blocked Countries" 3

    report "[$(format_date)] Retrieving data" 3
    retrieve_data
    report "[$(format_date)] DONE" 0
    echo ""

    report "[$(format_date)] Updating the chains" 3
    add_chain
    report "[$(format_date)] DONE" 0
    echo ""

    report "[$(format_date)] Cleaning up the unused ipsets and iptables rules" 3
    clean_old_relations
    report "[$(format_date)] DONE" 0
    echo ""

    report "[$(format_date)] Creating default whitelist" 3
    report "Kolmisoft access: 77.240.248.90 AND 5.9.57.125, API for ipdeny.com: 192.241.240.22, DNS, Local and M2 Servers will be added!" 3
    add_default_whitelist
    report "[$(format_date)] DONE" 0
    echo ""

    report "[$(format_date)] Updating the ipsets and iptables" 3
    add_rule
    report "[$(format_date)] DONE" 0

    report "[$(format_date)] Script completed" 0

elif  [ "$1" == "RENEW" ]; then

    report "[$(format_date)] Renewing Blocked Countries" 3

    rm -f /usr/local/m2/ip_zones/*.zone

    remove_chain $BLACKLIST_CHAIN

    report "[$(format_date)] Restarting ipset" 3
    service ipset restart &>/dev/null

    report "[$(format_date)] Script completed" 0

else
    report "Bad option! Try UPDATE or RENEW" 1
fi