#! /bin/bash
# This is script collects blocked/whitelisted IPs from iptables, compares with records from database and insert differences into database
# v 1.3

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

# Do not add INPUT to whitelist_chain
whitelist_chains="M2-WHITELIST-GUI"
blacklist_chains="INPUT M2-BLOCKED-IP-FROM-GUI"
ipauth_whitelist_chain="M2-CONNECT-POINTS-WHITELIST"
ipauth_whitelist_ips=""
server_config_file="/etc/m2/system.conf"
server_id=0
iptables_dir=/tmp/gui_iptables

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

db_blocked_ips=${iptables_dir}/m2_db_blocked_ips.txt
iptables_blocked_ips=${iptables_dir}/m2_iptables_blocked_ips.txt
blocked_ips_diff=${iptables_dir}/m2_blocked_ips_diff.txt
blocked_ips_to_insert=${iptables_dir}/m2_blocked_ips_to_insert.txt
blocked_ips_to_delete=${iptables_dir}/m2_blocked_ips_to_delete.txt

db_whitelisted_ips=${iptables_dir}/m2_db_whitelisted_ips.txt
iptables_whitelisted_ips=${iptables_dir}/m2_iptables_whitelisted_ips.txt
whitelisted_ips_diff=${iptables_dir}/m2_whitelisted_ips_diff.txt
whitelisted_ips_to_insert=${iptables_dir}/m2_whitelisted_ips_to_insert.txt
whitelisted_ips_to_delete=${iptables_dir}/m2_whitelisted_ips_to_delete.txt

# MySQL insert/delete batch size (how many IPs per query)
sql_batch_size=100

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 $db_blocked_ips &> /dev/null
    rm -f $iptables_blocked_ips &> /dev/null
    rm -f $blocked_ips_diff &> /dev/null
    rm -f $blocked_ips_to_insert &> /dev/null
    rm -f $blocked_ips_to_delete &> /dev/null
    rm -f $db_whitelisted_ips &> /dev/null
    rm -f $iptables_whitelisted_ips &> /dev/null
    rm -f $whitelisted_ips_diff &> /dev/null
    rm -f $whitelisted_ips_to_insert &> /dev/null
    rm -f $whitelisted_ips_to_delete &> /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"
}

insert_ips_to_database()
{
    local input_file=$1
    local type=$2
    local unblock_value=0
    local state_unblock_value=0
    local sql_insert_statement="INSERT INTO blocked_ips (server_id, blocked_ip, chain, unblock) VALUES"
    local sql_insert=$sql_insert_statement
    local insert_count=0
    local ip_address
    local ip_chain

    if [ "$type" == "blocked" ]; then
        insert_unblock_value=0
        state_unblock_value=1
    else
        insert_unblock_value=3
        state_unblock_value=4
    fi

    while IFS= read -r line; do
        ip_address=$(echo "$line" | grep -Po "^.*;" | tr -d ";")
        ip_chain=$(echo "$line" | grep -Po ";.*" | sed 's/^.//')

        # Do not insert if IP is already in being unblocked/unwhitelisted state
        state_count=$(mysql_query "SELECT COUNT(*) FROM blocked_ips WHERE server_id = $server_id AND unblock = $state_unblock_value AND blocked_ip = '$ip_address' AND chain = '$ip_chain'")
        if [[ -n "$state_count" && "$state_count" != "0" ]]; then
            log_notice "Will not insert $ip_address into DB as this IP already scheduled to be unblocked/unwhitelisted"
            continue;
        fi

        sql_insert="$sql_insert ($server_id, '$ip_address', '$ip_chain', $insert_unblock_value),"
        insert_count=$((insert_count + 1))

        # Execute full batch
        if (( insert_count >= sql_batch_size )); then
            # Remove last character (which should be comma)
            sql_insert="${sql_insert/%,/}"
            mysql_query "$sql_insert"
            sql_insert=$sql_insert_statement
            insert_count=0
        fi
    done < "$input_file"

    # Execute partial batch
    if (( insert_count > 0 )); then
        # Remove last character (which should be comma)
        sql_insert="${sql_insert/%,/}"
        mysql_query "$sql_insert"
    fi
}

delete_ips_from_database()
{
    local input_file=$1
    local type=$2
    local delete_count=0
    local sql_delete_statement
    local sql_delete
    local ip_address

    if [ "$type" == "blocked" ]; then
        unblock_value=0
    else
        unblock_value=3
    fi

    sql_delete_statement="DELETE FROM blocked_ips WHERE server_id = $server_id AND unblock = $unblock_value AND blocked_ip IN ("
    sql_delete=$sql_delete_statement

    while IFS= read -r line; do
        ip_address=$(echo "$line" | grep -Po "^.*;" | tr -d ";")
        sql_delete="$sql_delete'$ip_address',"
        delete_count=$((delete_count + 1))

        # Execute full batch
        if (( delete_count >= sql_batch_size )); then
            # Replace last character (which should be comma) with )
            sql_delete="${sql_delete/%,/)}"
            mysql_query "$sql_delete"
            sql_delete=$sql_delete_statement
            delete_count=0
        fi
    done < "$input_file"

    # Execute partial batch
    if (( delete_count > 0 )); then
        # Replace last character (which should be comma) with )
        sql_delete="${sql_delete/%,/)}"
        mysql_query "$sql_delete"
    fi
}

exist_in_whitelist()
{
    local ip_to_check="$1"
    local whitelist_ip

    for whitelist_ip in $ipauth_whitelist_ips; do
        if [ "$ip_to_check" == "$whitelist_ip" ]; then
            return 0
        fi
    done

    return 1
}


################## MAIN ##################


# 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 get blocked 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

# Give time to fail2ban to block all IPs from internal fail2abn DB
# This way we will not need to delete and imports IPs again on MOR DB
uptime_seconds=$(awk '{printf "%0.f", $1}' /proc/uptime)
if [[ -n $uptime_seconds ]] && ((uptime_seconds < 1800)); then
    if [[ -e /var/lib/fail2ban/fail2ban.sqlite3 ]]; then
        bips=$(sqlite3 /var/lib/fail2ban/fail2ban.sqlite3 "select count(*) from bips")

        if [[ -n $bips ]]; then
           # Give 10 seconds for 100 bips
           time_to_wait=$((bips / 10))
           if ((uptime_seconds < time_to_wait)); then
              log_notice "Waiting for $((time_to_wait - uptime_seconds)) seconds for fail2ban to load the IPs"
              exit 0
           fi
        fi
    fi
fi

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

# Get fail2ban chains
blacklist_chains+=" $(grep -F 'INPUT' $iptables_output | grep -Po '(f2b|fail2ban)-[\d\w-_]+' | sort -u)"

# Get whitelisted IPs from iptables
for chain in $whitelist_chains; do
    if grep -Fq -- "-A $chain" $iptables_output; then
        # With comments
        grep -- "-A $chain .*-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

        # Without comments
        grep -- "-A $chain .*-s.*ACCEPT" $iptables_output | grep -v -- "--comment" | sed -rn "s|(.*)-s ([[:digit:].]+/[[:digit:]]+)(.*)|\2;$chain|p" >> $iptables_whitelisted_ips
    fi
done

# Get blocked IPs from iptables
for chain in $blacklist_chains; do
    if grep -Fq -- "-A $chain" $iptables_output; then
        # With comments
        if [[ $chain == M2-BLOCKED-IP-FROM-GUI ]]; then
            grep -P -- "-A $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
        fi
        # Without comments
        grep -P -- "-A $chain .*-s.*(DROP|REJECT)" $iptables_output | grep -v -- "--comment" | sed -rn "s|(.*)-s ([[:digit:].]+/[[:digit:]]+)(.*)|\2;$chain|p" >> $iptables_blocked_ips
    fi
done

# Get whitelisted IP auth device IPs
ipauth_whitelist_ips=$(grep -- "-A $ipauth_whitelist_chain " $iptables_output | grep -F 'ACCEPT' | grep -Po '\d+\.\d+\.\d+\.\d+')

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

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

# Remove /32 from blocked IPs
sed -ri 's|([[:digit:].]+)/32|\1|' $iptables_blocked_ips

# Sort blocked IPs
sort -u $db_blocked_ips -o $db_blocked_ips
sort -u $iptables_blocked_ips -o $iptables_blocked_ips

# Remove /32 from whitelisted IPs
sed -ri 's|([[:digit:].]+)/32|\1|' $iptables_whitelisted_ips

# # Sort whitelisted IPs
sort -u $db_whitelisted_ips -o $db_whitelisted_ips
sort -u $iptables_whitelisted_ips -o $iptables_whitelisted_ips

# Find differences between blocked IPs in iptables and database
diff $iptables_blocked_ips $db_blocked_ips > $blocked_ips_diff
grep "^<" $blocked_ips_diff > $blocked_ips_to_insert
grep "^>" $blocked_ips_diff > $blocked_ips_to_delete

# # Find differences between whitelisted IPs in iptables and database
diff $iptables_whitelisted_ips $db_whitelisted_ips > $whitelisted_ips_diff
grep "^<" $whitelisted_ips_diff > $whitelisted_ips_to_insert
grep "^>" $whitelisted_ips_diff > $whitelisted_ips_to_delete

# Remove < > symbols from files
sed -i "s|< ||g" $blocked_ips_to_insert
sed -i "s|< ||g" $whitelisted_ips_to_insert
sed -i "s|> ||g" $blocked_ips_to_delete
sed -i "s|> ||g" $whitelisted_ips_to_delete

# Remove 0.0.0.0
sed -i "/;0.0.0.0/d" $blocked_ips_to_insert
sed -i "/;0.0.0.0/d" $whitelisted_ips_to_insert

# TODO: Fix performance when there are many devices
# Do not insert IP if it is present in IP auth device whitelist
while IFS=';' read -r ip_address; do
    if exist_in_whitelist "$ip_address"; then
        sed -i "/$ip_address/d" $blocked_ips_to_insert
   fi
done < <(cat $blocked_ips_to_insert | awk -F';' '{print $1}')

# Remove duplicates from blocked IPs
sort -u $blocked_ips_to_insert -o $blocked_ips_to_insert
sort -u $blocked_ips_to_delete -o $blocked_ips_to_delete

# Remove duplicates from whitelisted IPs
sort -u $whitelisted_ips_to_insert -o $whitelisted_ips_to_insert
sort -u $whitelisted_ips_to_delete -o $whitelisted_ips_to_delete

# Insert blocked IPs
if [ -s $blocked_ips_to_insert ]; then
    log_notice "Blocked IP list to insert into database"
    cat $blocked_ips_to_insert
    echo ""
    insert_ips_to_database $blocked_ips_to_insert "blocked"
fi

# Delete blocked IPs
if [ -s $blocked_ips_to_delete ]; then
    log_notice "Blocked IP list to delete from database"
    cat $blocked_ips_to_delete
    echo ""
    delete_ips_from_database $blocked_ips_to_delete "blocked"
fi

# Insert whitelisted IPs
if [ -s $whitelisted_ips_to_insert ]; then
    log_notice "Whitelisted IP list to insert into database"
    cat $whitelisted_ips_to_insert
    echo ""
    insert_ips_to_database $whitelisted_ips_to_insert "whitelisted"
fi

# Delete whitelisted IPs
if [ -s $whitelisted_ips_to_delete ]; then
    log_notice "Whitelisted IP list to delete from database"
    cat $whitelisted_ips_to_delete
    echo ""
    delete_ips_from_database $whitelisted_ips_to_delete "whitelisted"
fi

if [ -s $blocked_ips_to_insert ] || [ -s $blocked_ips_to_delete ] || [ -s $whitelisted_ips_to_insert ] || [ -s $whitelisted_ips_to_delete ]; then
    log_notice "Script finished"
else
    log_notice "No changes present, script finished"
fi

# Remove tmp files
remove_tmp_files