#!/bin/bash
# vim: set tabstop=4 shiftwidth=4 expandtab:
#
# This script up and stop SLB routes for default interface.
#
# Relevant variables for /etc/rc.conf.local is next:
#   ya_slb_enable   - if set to 'YES' add IPv4 SLB route
#   ya_slb6_enable  - if set to 'YES' add IPv6 SLB route
#   ya_slb_mtu      - MTU for SLB routes (default 8910)
#
# So, you may also specify this options in '/etc/network/interfaces' for choosen interface.
# You should note, that specifed interfaces must be a default gateway.
# For example:
#   iface eth2 inet6 auto
#       privext 0
#       mtu 8950
#       ya-slb yes              # the same as 'ya_slb_enable'
#       ya-slb6 yes             # the same as 'ya_slb6_enable'
#       ya-slb-mtu 8910         # the same as 'ya_slb_mtu'
#

#
# ya_slb_init_vars
#   Init basic variables.
#
ya_slb_init_vars()
{
    # constants
    YA_SLB_RUNDIR="/run/network/ya-slb"
    YA_SLB_CACHEDIR="/var/cache/network/ya-slb"
    YA_SLB_ROUTE_TABLE_IPV4="2010"
    YA_SLB_ROUTE_TABLE_IPV6="3010"

    local _config_file="/etc/rc.conf.local"

    # use configuration variables as local
    # for backward compatibility with /etc/rc.conf.local (ya.subr)
    local ya_slb_enable="NO"
    local ya_slb6_enable="NO"
    local ya_slb_mtu="8910"

    if [ -f ${_config_file} ] ; then
        . ${_config_file}
    fi

    # this variables will be exported
    : ${IF_YA_SLB:=$ya_slb_enable}
    : ${IF_YA_SLB6:=$ya_slb6_enable}
    : ${IF_YA_SLB_MTU:=$ya_slb_mtu}
    : ${IF_YA_SLB_CONF:="${YA_SLB_RUNDIR}/slb4.conf"}
    : ${IF_YA_SLB6_CONF:="${YA_SLB_RUNDIR}/slb6.conf"}
    : ${IF_YA_SLB_CONF_COMMON:="${YA_SLB_RUNDIR}/slb.conf"}
    : ${IF_YA_SLB_STATE:="${YA_SLB_RUNDIR}/state"}
}

#
# ya_slb_usage
#   Print usage information.
#
ya_slb_usage()
{
    echo "$(basename $0) start|stop|restart [interface]" >&2
}

#
# ya_check_enabled <variable_name>
#   Check that variable exists and it's value is 'yes', 'true', '1' or 'on'.
#   Return 0 if enabled, nonzero otherwise.
#
ya_check_enabled()
{
    local _value
    eval _value=\$${1}
    case ${_value} in
    # "yes", "true", "on", or "1"
    [Yy][Ee][Ss]|[Tt][Rr][Uu][Ee]|[Oo][Nn]|1)
        return 0
        ;;
    *)
        return 1
        ;;
    esac
}

#
# ya_log <message>
#   Print log message.
#
ya_log()
{
    echo "[$(date +%Y-%m-%d\ %H:%M:%S)] $@."
}

#
# ya_log_debug <message>
#   Print debug message.
#
ya_log_debug()
{
    if ya_check_enabled 'DEBUG' || ya_check_enabled 'IF_DEBUG' ; then
        ya_log "$@"
    fi
}

#
# ya_check_real_interface <interface>
#   Checks that interface is not loopback, tap, tun or other additional interface.
#
ya_check_real_interface()
{
    local _iface
    if [ ${#} -lt 1 -o -z "${1}" ] ; then
        ya_log_debug "ya_check_real_interface: no interface specified"
        return 2
    fi

    _iface=${1}
    # delete trailing numbers and downcase letters
    _iface="$(echo ${_iface%%[0-9]*} | tr '[:upper:]' '[:lower:]')"

    if [ -z ${_iface} ] ; then
         return 1
     fi

    case ${_iface} in
    lo|vlan|tun|ip6tnl|tap|vif|ipfw|pflog|virbr|plip|dummy)
        return 1
        ;;
    --all|all)
        return 1
        ;;
    esac

    return 0
}

#
# ya_get_active_interface <interface>
#   Get active interface:
#       - interface must be as a default gateway.
#       - interface must has a real name.
#
ya_get_active_interface()
{
    local _ifaces _iface
    _iface=${1:-''}

    _ifaces=""
    _ifaces=$(ip -4 route list exact 0.0.0.0/0 scope global | grep -v 'dev tun' | grep -oP 'dev\s+\K\w+')
    _ifaces="$_ifaces $(ip -6 route list exact ::/0 scope global | grep -oP 'dev\s+\K\w+')"
    _ifaces=$(echo $_ifaces | tr " " "\n" | uniq)

    if [ ! -z "$_iface" ] ; then
        _ifaces=$(echo $_ifaces | grep -o $_iface)
    fi

    for _iface in $_ifaces ; do
        if ya_check_real_interface $_iface ; then
            echo $_iface
            return
        fi
    done

    echo -n ''
    return
}

#
# ya_get_ip4addr <iface>
#   Print global IPv4 address of specified interface.
#   If address not found or interface specifed 'none' will be returned.
#
ya_get_ip4addr()
{
    local _iface _ip

    if [ ${#} -ne 1 ] ; then
        echo 'none'
        return 1
    fi

    _iface=${1}
    _ip=$(ip -o -4 addr show scope global dev ${_iface} |\
        awk '! / (10\.|192\.168).*/ {sub(/\/.*/, "", $4); print $4; exit}')

    if [ -z "${_ip}" ] ; then
        echo 'none'
        return 2
    fi

    echo "${_ip}"
    return 0
}

#
# ya_get_ip6addr <iface>
#   Print global IPv6 address of specified interface.
#   If address not found or interface specifed 'none' will be returned.
#
ya_get_ip6addr()
{
    local _iface _ip

    if [ ${#} -ne 1 ] ; then
        echo 'none'
        return 1
    fi

    _iface=${1}
    _ip=$(ip -o -6 addr show scope global dev ${_iface} |\
            awk '! / (fd|fc).*/ && ! /temporary/ {sub(/\/.*/, "", $4); print $4; exit}')

    if [ -z "${_ip}" ] ; then
        echo "none"
        return 2
    fi

    echo "${_ip}"
    return 0
}

#
# ya_is_ipv4 <ip>
#   Returns 0 if 'ip' is IPv4 address, otherwise returns 1.
#
ya_is_ipv4()
{
    if [ $# -ne 1 -o -z "$1" ] ; then
        return 1
    fi

    if [ -z "${1##*.*}" ] ; then
        return 0
    fi

    return 1
}

#
# ya_is_ipv6 <ip>
#   Returns 0 if 'ip' is IPv6 address, otherwise returns 1.
#
ya_is_ipv6()
{
    if [ $# -ne 1 -o -z "$1" ] ; then
        return 1
    fi

    if [ -z "${1##*:*}" ] ; then
        return 0
    fi

    return 1
}

#
# ya_get_cidr <ip|cidr>
#   Print CIDR notation for specified IP address or CIDR.
#   This function use 24 netmask for IPv4 and 64 for IPv6 if prefix not specified.
#
ya_get_cidr()
{
    local _cidr _ip
    if [ ${#} -ne 1 -o -z "${1}" ] ; then
        ya_log_debug "ya_get_cidr params error: no IP/CIDR specified"
        return 1
    fi

    _cidr=${1}
    _ip=${_cidr%/*}

    # already in cidr
    if [ "${_ip}" != "${_cidr}" ] ; then
        echo $_cidr
        return
    fi

    if ya_is_ipv4 ${_ip} ; then
        echo "${_ip}/24"
    fi

    if ya_is_ipv6 ${_ip} ; then
        echo "${_ip}/64"
    fi

    echo -n

    return
}

#
# ya_get_ipv4_max_host <cidr>
#   Print max IPv4 host in specified network.
#
ya_get_ipv4_max_host()
{
    #
    # ya_ipv4_prefix2mask prefix
    #   Convert cidr prefix to ipv4 netmask.
    #
    ya_ipv4_prefix2mask()
    {
        set -- $(( 5 - ($1 / 8) )) 255 255 255 255 \
            $(( (255 << (8 - ($1 % 8))) & 255 )) 0 0 0
        [ ${1} -gt 1 ] && shift $1 || shift
        echo ${1-0} ${2-0} ${3-0} ${4-0}
    }

    local _cidr _ipaddr _prefix
    local o1 o2 o3 o4 m1 m2 m3 m4 ip1 ip2 ip3 ip4

    _cidr=${1}
    _ipaddr=${_cidr%/*}
    _prefix=${_cidr#*/}
    echo ${_ipaddr} | sed -e 's/\./ /g' | while read o1 o2 o3 o4 ; do
        ya_ipv4_prefix2mask ${_prefix} | while read m1 m2 m3 m4 ; do
            ip1=$(($o1 + (255 - ($o1 | $m1))))
            ip2=$(($o2 + (255 - ($o2 | $m2))))
            ip3=$(($o3 + (255 - ($o3 | $m3))))
            ip4=$(($o4 + (255 - ($o4 | $m4))))
            if [ ${_prefix} -lt 31 ] ; then
                ip4=$((${ip4} - 1))
            fi
            echo "${ip1}.${ip2}.${ip3}.${ip4}"
        done
    done

}

#
# ya_slb_is_configured
#   Check that SLB tun already configured.
#
ya_slb_is_configured()
{
    if [ -f ${IF_YA_SLB_STATE} ] ; then
        cat ${IF_YA_SLB_STATE}
        return 0
    fi

    echo -n ''
    return 1
}

#
# ya_slb_mark_configured
#   Set a flag, that ya-slb configured.
#
ya_slb_mark_configured()
{
    echo $@ > ${IF_YA_SLB_STATE}

    return 0
}

#
# ya_slb_unmark_configured
#   Unset a flag, that ya-slb configured.
#
ya_slb_unmark_configured()
{
    rm -f ${IF_YA_SLB_STATE}

    return 0
}

#
# ya_slb_fetcher <ip> <conf>
#   Fetch SLB configuragtion from Racktables API by IP to file.
#
ya_slb_fetcher()
{
    local _ip _conf _fetch_args _wget_args _url _fetch_cmd _max_tries _cache

    if [ ${#} -ne 2 -o -z "${1}" -o -z "${2}" ] ; then
        ya_log_debug "ya_slb_fetcher params error: 2 params required"
        return 1
    fi

    # initialize variables
    _ip="${1}"
    _conf="${2}"
    _cache="${YA_SLB_CACHEDIR}/$(basename ${_conf})"
    _url="https://ro.racktables.yandex.net/export/slb-info.php?mode=aliases&ip=${_ip}"
    _fetch_args="-q -T 15 -o ${_conf}"
    _wget_args="-q -t 1 -T 15 -O ${_conf} --no-check-certificate"
    _max_tries=4
    if ! _fetch_cmd="$(which fetch) ${_fetch_args} ${_url}" ; then
        if ! _fetch_cmd="$(which wget) ${_wget_args} ${_url}" ; then
            ya_log "Can't find fetching command (fetch and wget doesn't exists)"
            return 3
        fi
    fi

    # make empty file
    echo -n > ${_conf}
    chmod 644 ${_conf}

    mkdir -p $(dirname ${_cache})

    # check that IP is specified
    if [ -z "$_ip" -o "$_ip" = "none" ] ; then
        ya_log_debug "IP address for fetching SLB TUN information not found"
        return 2
    fi

    # try fetch file
    local i=0
    ya_log_debug "Fetching SLB configuration ${_fetch_cmd}"
    while [ ! ${i} -eq "${_max_tries}" ] ; do
        if ${_fetch_cmd} > /dev/null 2>&1 ; then
            cp ${_conf} ${_cache}
            return 0
        fi

        i=$((${i} + 1))
        sleep ${i}
    done

    # if we can't fetch a file, we may use cached version.
    if [ -f ${_cache} ] ; then
        cp ${_cache} ${_conf}
    fi

    return 4
}

#
# ya_slb_get_config <interface> [<config_name>]
#   Get common SLB configuration for <interface> and save it to <config_name>.
#
ya_slb_get_config()
{
    local _iface _conf _ip4 _ip6

    if [ ${#} -lt 1 ] ; then
        ya_log_debug "ya_slb_get_config expected minimum 1 params: '${#}' given"
        return 1
    fi

    _iface=${1}
    _conf=${2:-$IF_YA_SLB_CONF_COMMON}
    _ip4=$(ya_get_ip4addr ${_iface})
    _ip6=$(ya_get_ip6addr ${_iface})

    if ya_check_enabled 'IF_YA_SLB' ; then
        ya_slb_fetcher ${_ip4} ${IF_YA_SLB_CONF}
    else
        echo -n > ${IF_YA_SLB_CONF}
    fi

    if ya_check_enabled 'IF_YA_SLB6' ; then
        ya_slb_fetcher ${_ip6} ${IF_YA_SLB6_CONF}
    else
        echo -n > ${IF_YA_SLB6_CONF}
    fi

    cat ${IF_YA_SLB_CONF} ${IF_YA_SLB6_CONF} | sort | uniq > ${_conf}
    chmod 644 ${_conf}

    if [ "$(wc -l ${_conf})" != "0" ] ; then
        return 0
    fi

    return 2
}

#
# ya_slb_get_slb_info <slb_ip>
#   Print slb_* options by specified SLB IP address (cidr)
#
#
ya_slb_get_slb_info()
{
    local _cidr _ip _proto _table _router _mtu _mss
    _cidr=$(ya_get_cidr ${1})

    _ip=${_cidr%/*}
    _proto=""
    _table=""
    _router=""
    _mtu=${IF_YA_SLB_MTU}
    _mss=""

    if ya_is_ipv4 ${_ip} ; then
        _proto="4"
        _table=${YA_SLB_ROUTE_TABLE_IPV4}
        _router=$(ya_get_ipv4_max_host ${_cidr})
        _mss=$((${_mtu} - 40))
    elif ya_is_ipv6 ${_ip} ; then
        _proto="6"
        _table=${YA_SLB_ROUTE_TABLE_IPV6}
        _router="$(echo ${_ip} | sed -r '/(^|\s)([A-Fa-f0-9]+:){4}/! {s/::/:0:/}' | grep -Eoh '([0-9a-f]+:){4}'):1"
        _mss=$((${_mtu} - 60))
    fi

    echo "slb_proto=${_proto} slb_cidr=${_cidr} slb_ip=${_ip} slb_table=${_table} slb_router=${_router} slb_mss=${_mss} slb_mtu=${_mtu}"
}

#
# ya_slb_start_slb <interface> <config>
#   Up SLB address.
#
ya_slb_start_slb()
{
    local _iface _config _slb_ip _rule _ip
    local slb_proto slb_cidr slb_router slb_mtu slb_mss slb_table

    if [ ${#} -ne 2 -o -z "${1}" -o -z "${2}" ] ; then
        ya_log_debug "ya_slb_start_slb params error: 2 params required"
        return 1
    fi

    _rule=0
    _iface=${1}
    _config=${2}
    while read _slb_ip ; do
        eval `ya_slb_get_slb_info ${_slb_ip}`
        if [ -z "${slb_proto}" ] ; then
            ya_log "Error record '${_slb_ip}' record in ${_config}. Unknown IP."
            continue
        fi
        _table=$((${slb_table} + ${_rule}))

        ya_log "Add IP ${slb_cidr} to ${_iface}"
        ip -${slb_proto} addr add ${slb_cidr} dev ${_iface}
        ya_log "Add rule from ${slb_ip} to ${_table}"
        ip -${slb_proto} rule add from ${slb_ip} lookup ${_table} priority ${_table}
        ya_log "Add default route via ${slb_router} from table ${_table}"
        ip -${slb_proto} route add default via ${slb_router} dev ${_iface} \
            table ${_table} mtu ${slb_mtu} advmss ${slb_mss}

        _rule=$((${_rule} + 1))
    done < ${_config}
}

#
# ya_slb_stop_slb <interface> <config>
#   Down SLB.
#
ya_slb_stop_slb()
{
    local _iface _config _slb_ip _rule _ip
    local slb_proto slb_cidr slb_router slb_mtu slb_mss slb_table

    if [ ${#} -ne 2 -o -z "${1}" -o -z "${2}" ] ; then
        ya_log_debug "ya_slb_stop_slb params error: 2 params required"
        return 1
    fi

    _rule=0
    _iface=${1}
    _config=${2}
    while read _slb_ip ; do
        eval `ya_slb_get_slb_info ${_slb_ip}`
        if [ -z "${slb_proto}" ] ; then
            ya_log "Error record '${_slb_ip}' record in ${_config}. Unknown IP"
            continue
        fi
        _table=$((${slb_table} + ${_rule}))

        ya_log "Remove rule ${_table} lookup ${slb_ip}"
        ip -${slb_proto} rule del from ${slb_ip} lookup ${_table}
        ya_log "Remove default route from table ${_table}"
        ip -${slb_proto} route del default via ${slb_router} dev ${_iface} \
            table ${_table}
        ya_log "Remove IP address from ${slb_cidr} dev ${_iface}"
        ip -${slb_proto} addr del ${slb_cidr} dev ${_iface}

        _rule=$((${_rule} + 1))
    done < ${_config}
}

#
# ya_slb_run_start [<interface>]
#   Start SLB.
#
ya_slb_run_start()
{
    local _iface _state _dummy_iface _config

    if [ ! -z "$(ya_slb_is_configured)" ] ; then
        ya_log_debug "SLB already configured"
        return 0
    fi

    _iface=$(ya_get_active_interface ${1:-${IFACE:-''}})
    if [ -z "${_iface}" ] ; then
        ya_log_debug "No active interface found or given interface not supported"
        return 1
    fi

    _config=${IF_YA_SLB_CONF_COMMON}
    # get SLB configuration
    if ! ya_slb_get_config ${_iface} ${_config} ; then
        ya_log_debug "No configuration found or fetching failed for '${_iface}'"
        return 1
    fi

    if [ ! -s ${_config} ]; then
        ya_log "Empty SLB config for ${_iface}"
        return 0
    fi

    ya_slb_start_slb ${_iface} ${_config}

    ya_slb_mark_configured ${_iface}

    return 0
}

#
# ya_slb_run_stop [<interface>]
#   Stop SLB tunnel.
#
ya_slb_run_stop()
{
    local _iface _state _config

    _state=$(ya_slb_is_configured)
    _iface=${1:-${IFACE:-${_state:-''}}}
    _config=${IF_YA_SLB_CONF_COMMON}
    if [ -z "${_state}" ] ; then
        ya_log_debug "ya-slb not configured. Nothing to stop"
        return 0
    fi

    if [ "${_state}" != "${_iface}" ] ; then
        ya_log_debug "Specified interface '${_iface}', but configured '${_state}'. Stop skipping"
        return 0
    fi

    if [ ! -e ${_config} ] ; then
        if ! ya_slb_get_config ${_iface} ${_config} ; then
            ya_log_debug "No configuration found or fetching failed"
            return 0
        fi
    fi

    ya_slb_stop_slb ${_iface} ${_config}
    ya_slb_unmark_configured

    return 0
}

#
# ya_slb_run_restart [<interface>]
#   Start SLB tunnel.
ya_slb_run_restart()
{
    local _iface

    _iface=${1:-''}
    ya_slb_run_stop ${_iface}
    sleep 2;
    ya_slb_run_start ${_iface}

    return 0
}

#
# ya_slb_run <action> [<interface>]
#   Run a script.
#
ya_slb_run()
{
    local _action=${1:-${MODE:-"none"}}
    local _iface=${2:-''}

    ya_slb_init_vars
    mkdir -p ${YA_SLB_RUNDIR} || exit 0

    if ! (ya_check_enabled 'IF_YA_SLB' || ya_check_enabled 'IF_YA_SLB6') ; then
        ya_log_debug "$(basename $0) is off"
        exit 0
    fi

    case "${_action}" in
    start)
        ya_log_debug "Start SLB"
        ya_slb_run_start ${_iface}
        exit 0
        ;;
    stop)
        ya_log_debug "Stop SLB"
        ya_slb_run_stop ${_iface}
        exit 0
        ;;
    restart)
        ya_log_debug "Restart SLB"
        ya_slb_run_restart ${_iface}
        exit 0
        ;;
    *)
        ya_slb_usage
        exit 0
        ;;
    esac
}

ya_slb_run "$@"
