#!/bin/bash

# Kolmisoft, 2016-2023
# Script check the database and block/whitelist selected IPs
# v 1.4.0

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

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

server_config_file="/etc/m2/system.conf"
block_command_output=""
whitelist_command_output=""
block_chain="M2-BLOCKED-IP-FROM-GUI"
whitelist_chain_gui="M2-WHITELIST-GUI"
whitelist_chain_conn_points="M2-CONNECT-POINTS-WHITELIST"
iptables_dir=/tmp/gui_iptables/m2_block_ip

# Temporary files
iptables_output=${iptables_dir}/m2_iptables_output.txt

iptables_blocked_ips=${iptables_dir}/m2_iptables_blocked_ips.txt
db_ips_to_block=${iptables_dir}/m2_db_ips_to_block.txt
ips_to_block_diff=${iptables_dir}/m2_to_block_diff.txt

db_ips_to_whitelist=${iptables_dir}/m2_db_ips_to_whitelist.txt
iptables_whitelisted_ips=${iptables_dir}/m2_iptables_whitelisted_ips.txt
ips_to_whitelist_diff=${iptables_dir}/m2_ips_to_whitelist_diff

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

log_notice() {
    echo "[$(format_date)] [NOTICE] $*"
}

log_error() {
    echo "[$(format_date)] [ERROR] $*"
}

remove_tmp_files()
{
    rm -f $iptables_output &> /dev/null
    rm -f $iptables_blocked_ips &> /dev/null
    rm -f $db_ips_to_block &> /dev/null
    rm -f $ips_to_block_diff &> /dev/null
    rm -f $iptables_whitelisted_ips &> /dev/null
    rm -f $db_ips_to_whitelist &> /dev/null
    rm -f $ips_to_whitelist_diff &> /dev/null
}

mysql_query() {
    local query="$1"
    local result

    result=$(MYSQL_PWD="$DB_PASSWORD" /usr/bin/mysql -h "$DB_HOST" -u "$DB_USERNAME" $P_OPT "$DB_NAME" -sNe "$query")
    echo -n "$result"
}

block_command() {
    /sbin/iptables $l_opt -I $block_chain -s "$ip_address" -j DROP -m comment --comment "$comment"
}

whitelist_command() {
    /sbin/iptables $l_opt -I $whitelist_chain_gui -s "$ip_address" -j ACCEPT -m comment --comment "$comment"
}

fix_gui_ip_whitelist_priority() {
    # M2-WHITELIST-GUI chain should always be below M2-CONNECT-POINTS-WHITELIST

    # Check which priority has chain M2-CONNECT-POINTS-WHITELIST
    whitelist_chain_conn_points_priority=$(/sbin/iptables $l_opt -L -n --line-numbers | grep $whitelist_chain_conn_points | grep -v "Chain" | head -n 1 | awk '{print $1}')

    # Check which priority has chain M2-WHITELIST-GUI
    whitelist_chain_gui_priority=$(/sbin/iptables $l_opt -L -n --line-numbers | grep $whitelist_chain_gui | grep -v "Chain" | head -n 1 | awk '{print $1}')

    # If we both chains are added to input and we got their priorities, check if order is correct
    if [[ $whitelist_chain_conn_points_priority =~ ^[0-9]+$ ]]; then
        # Check if GUI IP whitelist chain priority is lower than connection points whitelist chain priority (by 1)
        expected_gui_chain_priprity=$(( whitelist_chain_conn_points_priority + 1 ))

        if [ "$whitelist_chain_gui_priority" != "$expected_gui_chain_priprity" ]; then
            report "Changing chain [$whitelist_chain_gui] priority from [$whitelist_chain_gui_priority] to [$expected_gui_chain_priprity]" 3

            /sbin/iptables $l_opt -D INPUT -j $whitelist_chain_gui
            /sbin/iptables $l_opt -I INPUT $expected_gui_chain_priprity -j $whitelist_chain_gui
        else
            # report "Chain [$whitelist_chain_gui] priority [$whitelist_chain_gui_priority] is ok" 0
            :
        fi
    else
        # Something is wrong... just set GUI IP whitelist chain priority to 1
        report "Failed to determine priority for chain [$whitelist_chain_conn_points], got priority [$whitelist_chain_conn_points_priority]" 1

        if [ "$whitelist_chain_gui_priority" != "1" ]; then
            report "Priority for chain [$whitelist_chain_gui] will be set to 1" 2
            /sbin/iptables $l_opt -D INPUT -j $whitelist_chain_gui
            /sbin/iptables $l_opt -I INPUT 1 -j $whitelist_chain_gui
        fi
    fi
}

# Protection against running more than one instance
if [[ "$(/sbin/pidof -x "$(basename "$0")" -o %PPID)" ]]; then
    log_notice "$(basename "$0") script is already running with PID $(/sbin/pidof -x "$(basename "$0")" -o %PPID)"
    exit 0
fi

log_notice "Starting M2 block/whitelist IP script"

set_database_variables

# Create dir for M2 GUI iptables
mkdir -p $iptables_dir

# Remove tmp files
remove_tmp_files

# Check if GUI iptables are enabled
m2_gui_iptables=$(sed 's/ //g' $server_config_file | awk -F"=" '/m2_gui_iptables/{print $2}')
if [ -n "$m2_gui_iptables" ] && [ "$m2_gui_iptables" -eq "0" ]; then
    log_notice "M2 GUI iptables is disabled. Exiting."
    exit 0
fi

# Get server id
server_id=$(grep server_id $server_config_file | tr -d '[:space:]' | awk -F'=' '{print $2}')
if [ "$server_id" -eq "0" ] || ! echo "$server_id" | grep -qE '^[0-9]+$'; then
    log_error "server_id! It has value $server_id, but it should be a digit > 0"
    log_error "Aborting script. Check server_id value in $server_config_file"
    exit 1
fi

# Save iptables output which will be used for parsing (instead of calling iptables command each time)
/sbin/iptables-save > $iptables_output


# create "M2-BLOCKED-IP-FROM-GUI" chain if it does not exist
if ! grep -Fv INPUT $iptables_output | grep -Fq "$block_chain"; then
    /sbin/iptables $l_opt -N "$block_chain" 2>/dev/null 1>&2

fi

# If chain is not referenced in INPUT, add it to INPUT
if ! grep -F INPUT $iptables_output | grep -Fq "$block_chain"; then
    /sbin/iptables $l_opt -A INPUT -j "$block_chain"
fi

# If chain does not exist, create one
if ! grep -Fv INPUT $iptables_output | grep -Fq "$whitelist_chain_gui"; then
    /sbin/iptables $l_opt -N "$whitelist_chain_gui" 2>/dev/null 1>&2
     # Add RETURN statement (not realy do anything, just for cosmetics)
    /sbin/iptables $l_opt -A "$whitelist_chain_gui" -j RETURN
fi

# If chain is not referenced in INPUT, add it to INPUT
if ! grep -F INPUT $iptables_output | grep -Fq "$whitelist_chain_gui"; then
    /sbin/iptables $l_opt -I INPUT -j "$whitelist_chain_gui"
fi

# Make sure that GUI IP whitelist priority is correct
fix_gui_ip_whitelist_priority

# Get blocked IPs from DB
mysql_query "SELECT CONCAT(blocked_ip, ';', chain) FROM blocked_ips WHERE server_id = $server_id and unblock = 2" > $db_ips_to_block

# Add \32 for single IPs
sed -ri 's|(\.[[:digit:]]+);|\1/32;|' $db_ips_to_block
# sort and leave unique only
sort -u $db_ips_to_block -o $db_ips_to_block

# Get blocked IPs from DB from blacklist chain
grep -P -- "-A $block_chain .*-s.*--comment.*(DROP|REJECT)" $iptables_output | perl -ne '/-s\s(\d+\.\d+\.\d+\.\d+\/\d+).+--comment\s(?:["]([^"]+)["]|([^\s"]+)\s)/ && print "$1;$2$3\n"' >> $iptables_blocked_ips
grep -P -- "-A $block_chain .*-s.*(DROP|REJECT)" $iptables_output | grep -v -- "--comment" | sed -rn "s|(.*)-s ([[:digit:].]+/[[:digit:]]+)(.*)|\2;|p" >> $iptables_blocked_ips
sort -u $iptables_blocked_ips -o $iptables_blocked_ips

# Find differences between blocked IPs database and iptab;es
grep "^<" < <(diff $db_ips_to_block $iptables_blocked_ips) | sed "s|< ||g" > $ips_to_block_diff

# Get whitelisted IPs from DB
mysql_query "SELECT CONCAT(blocked_ip, ';', chain) FROM blocked_ips WHERE server_id = $server_id and unblock = 5" > $db_ips_to_whitelist

# Add \32 for single IPs
sed -ri 's|(\.[[:digit:]]+);|\1/32;|' $db_ips_to_whitelist
# sort and leave unique only
sort -u $db_ips_to_whitelist -o $db_ips_to_whitelist

# Get whitelisted IPs from DB from blacklist chain
grep -- "-A $whitelist_chain_gui .*-s.*--comment.*ACCEPT" $iptables_output | perl -ne '/-s\s(\d+\.\d+\.\d+\.\d+\/\d+).+--comment\s(?:["]([^"]+)["]|([^\s"]+)\s)/ && print "$1;$2$3\n"' >> $iptables_whitelisted_ips
grep -- "-A $whitelist_chain_gui .*-s.*ACCEPT" $iptables_output | grep -v -- "--comment" | sed -rn "s|(.*)-s ([[:digit:].]+/[[:digit:]]+)(.*)|\2;|p" >> $iptables_whitelisted_ips
sort -u $iptables_whitelisted_ips -o $iptables_whitelisted_ips

# Find differences between whitelisted IPs in database and iptables
grep "^<" < <(diff $db_ips_to_whitelist $iptables_whitelisted_ips) | sed "s|< ||g" > $ips_to_whitelist_diff

if  [[ ! -s $ips_to_block_diff && ! -s $ips_to_whitelist_diff ]]; then
    log_notice "No changes present, script finished"
    remove_tmp_files
    exit 0
fi

while IFS=';' read -r ip_address comment; do
    if [[ -z $comment ]]; then
        if ! grep -Fq "$ip_address" $iptables_blocked_ips; then
            log_notice "Executing command /sbin/iptables -I $block_chain -s $ip_address -j DROP"
            if ! block_command_output=$(/sbin/iptables $l_opt -I $block_chain -s "$ip_address" -j DROP 2>&1); then
                iptables_error="1"
            fi
        fi
    else
        if ! grep -Fq "$ip_address" $iptables_blocked_ips; then
            log_notice "Executing command  /sbin/iptables -I $block_chain -s $ip_address -j DROP -m comment --comment \"$comment\""
            # do not change this
            # block_command has to be function
            # http://mywiki.wooledge.org/BashFAQ/050
            if ! block_command_output=$(block_command "$comment" 2>&1); then
                    iptables_error="1"
            fi

            if [ "$iptables_error" = "1" ]; then
                log_error "iptables returned such error:"
                echo "$block_command_output"
            fi
        fi
    fi
done < "$ips_to_block_diff"

while IFS=';' read -r ip_address comment; do
    if [[ -z $comment ]]; then
        if ! grep -Fq "$ip_address" $iptables_whitelisted_ips; then
            log_notice "Executing command /sbin/iptables -I $whitelist_chain_gui -s $ip_address -j ACCEPT"
            if ! whitelist_command_output=$(/sbin/iptables $l_opt -I $whitelist_chain_gui -s "$ip_address" -j ACCEPT 2>&1); then
                iptables_error="1"
            fi
        fi
    else
        if ! grep -Fq "$ip_address" $iptables_whitelisted_ips; then
            log_notice "Executing command /sbin/iptables -I $whitelist_chain_gui -s $ip_address -j ACCEPT -m comment --comment \"$comment\""
            # do not change this
            # block_command has to be function
            # http://mywiki.wooledge.org/BashFAQ/050
            if ! whitelist_command_output=$(whitelist_command "$comment" 2>&1); then
                    iptables_error="1"
            fi

            if [ "$iptables_error" = "1" ]; then
                log_error "iptables returned such error:"
                echo "$whitelist_command_output"
            fi
        fi
    fi
done < "$ips_to_whitelist_diff"

# Remove tmp files
remove_tmp_files