#!/bin/bash

SLB_MACRO_URL="https://racktables.yandex.net/export/expand-fw-macro.php?macro=_SLB_LO_NETS_"
SLB_CFG_DIR="/usr/share/manage-balancer/"
SLB_V4_CONFIG=$SLB_CFG_DIR"slb.v4"
SLB_V6_CONFIG=$SLB_CFG_DIR"slb.v6"

YANDEX_HFB_AGENT_DIR="/etc/yandex-hbf-agent/"

function gen_iptables_config() {
    echo "*filter"
    echo ":SLB -"
    echo "-A INPUT -j SLB"
    for network in $*; do 
        echo '-A SLB -p tcp -j REJECT -m comment --comment "Hide from balancer" -s '$network
    done
    echo "COMMIT"
}


function error_notify() {
    echo
    echo "===================================================="
    echo " Attention! Some problems with open/close this host"
    echo "===================================================="
}


function check_iptables() {
    iptables -L -n &>/dev/null
    if [ $? -ne 0 ];then
        echo "========================================================================="
        echo " Attention! Iptables disabled"
        echo " Turn on CAP_NET_ADMIN, CAP_NET_RAW features in qloud component settings"
        echo "========================================================================="
        return 1
    fi
}


function get_chain_list() {
    CHAIN_LIST=$( (iptables -S; ip6tables -S)  | grep "Hide from balancer" | awk '{print $2}' | sort | uniq )
    echo $CHAIN_LIST
}


function list() {
    check_iptables &&
    
    CHAIN_LIST=$(get_chain_list)
    if [ "$CHAIN_LIST" == "" ];then
        echo "No SLB rules"
        return 0
    fi

    for chain in $CHAIN_LIST; do
        echo "Chain "${chain}":"
        echo " "$(iptables -S $chain | grep -c "Hide from balancer") IPv4 rules
        echo " "$(ip6tables -S $chain | grep -c "Hide from balancer") IPv6 rules
    done
}


function update() {
    NETWORK_COUNT=20
    NETWORK_TMP_FILE="/tmp/slb_networks"
    curl -s $SLB_MACRO_URL > $NETWORK_TMP_FILE
    if [ $(wc -l $NETWORK_TMP_FILE | awk '{print $1}') -lt $NETWORK_COUNT ];then
        echo "Network list is too short!"
    fi
    
    IPv4_NETWORKS=`grep "\." $NETWORK_TMP_FILE`
    IPv6_NETWORKS=`grep ":" $NETWORK_TMP_FILE`

    gen_iptables_config $IPv4_NETWORKS > $SLB_V4_CONFIG
    gen_iptables_config $IPv6_NETWORKS > $SLB_V6_CONFIG

    echo -e "\n==================================== IPv4 CONFIG ===================================="
    cat $SLB_V4_CONFIG
    echo -e "\n==================================== IPv6 CONFIG ===================================="
    cat $SLB_V6_CONFIG
}


function create_custom_config() {
    CUSTOM_CHAIN=$1
    CUSTOM_V4_CONFIG=${SLB_CFG_DIR}${CUSTOM_CHAIN}".v4"
    CUSTOM_V6_CONFIG=${SLB_CFG_DIR}${CUSTOM_CHAIN}".v6"
    cp $SLB_V4_CONFIG $CUSTOM_V4_CONFIG &&
    cp $SLB_V6_CONFIG $CUSTOM_V6_CONFIG &&
    sed -i "s/SLB/$CUSTOM_CHAIN/" $CUSTOM_V4_CONFIG &&
    sed -i "s/SLB/$CUSTOM_CHAIN/" $CUSTOM_V6_CONFIG
}


function close() {
    check_iptables &&
    iptables-restore -n ${SLB_CFG_DIR}${CHAIN,,}".v4" &&
    ip6tables-restore -n ${SLB_CFG_DIR}${CHAIN,,}".v6" || error_notify
    if [ -d "$YANDEX_HFB_AGENT_DIR" ]; then
        ln -svf ${SLB_CFG_DIR}${CHAIN,,}".v4" $YANDEX_HFB_AGENT_DIR"rules.d/" &&
        ln -svf ${SLB_CFG_DIR}${CHAIN,,}".v6" $YANDEX_HFB_AGENT_DIR"rules.d/" || error_notify
    fi
    list
}


function open() {
    CHAIN=$1
    check_iptables &&
    if [ "$CHAIN" == "SLB" ];then
         for custom_chain in $(get_chain_list); do
            iptables -F $custom_chain &&
            ip6tables -F $custom_chain || error_notify

            rm -rfv $YANDEX_HFB_AGENT_DIR"rules.d/"$custom_chain".v4" &&
            rm -rfv $YANDEX_HFB_AGENT_DIR"rules.d/"$custom_chain".v6" || error_notify
         done
    else
        iptables -F $CHAIN &&
        ip6tables -F $CHAIN || error_notify
    fi

    if [ -d "$YANDEX_HFB_AGENT_DIR" ]; then
        rm -rfv $YANDEX_HFB_AGENT_DIR"rules.d/"${CHAIN,,}".v4" &&
        rm -rfv $YANDEX_HFB_AGENT_DIR"rules.d/"${CHAIN,,}".v6" || error_notify
    fi

    list
}

if [ -z $2 ];then
    CHAIN="SLB"
else
    CHAIN=$2
fi

case $1 in
    close)
    if [ "$CHAIN" != "SLB" ];then
        # Проверяем наличие в системе кастомной цепочки (RTEC-3549):
        if [ `echo $(get_chain_list)  | grep -c $CHAIN` -eq 0 ];then
            create_custom_config $CHAIN
        fi
    fi
    close $CHAIN
    ;;

    open)
    open $CHAIN
    ;;
    
    list|show)
    list
    ;;

    update)
    update
    ;;    
    *)  
        echo "Usage:"
        echo "manage-balancer [close|open|list|update]"
    ;;
esac

