#! /bin/bash
### BEGIN INIT INFO
# Provides:          nperf-server
# Required-Start:    $all
# Required-Stop:
# Default-Start:     2 3 4 5
# Default-Stop:
# Short-Description: nPerfServer
### END INIT INFO

# Function to determine the distribution
get_distro() {
    if [ -f /etc/redhat-release ]; then
        DISTRO=$(cat /etc/redhat-release | awk '{print $1}')
    elif [ -f /etc/lsb-release ]; then
        DISTRO=$(grep DISTRIB_ID /etc/lsb-release | awk -F= '{print $2}')
    elif [ -f /etc/debian_version ]; then
        DISTRO="Debian"
    elif [ -f /etc/SuSE-release ]; then
        DISTRO="SUSE"
    elif [ -f /etc/os-release ]; then
        DISTRO=$(grep ^ID= /etc/os-release | awk -F= '{print $2}')
    else
        DISTRO="Unknown"
    fi
    echo $DISTRO
}

# Get the distribution
DISTRO=$(get_distro)

# Required configuration to be present in the nftables configuration file
REQUIRED_CONF='#!/usr/sbin/nft -f

flush ruleset

table inet filter {
        chain input {
                type filter hook input priority 0;
        }
        chain forward {
                type filter hook forward priority 0;
        }
        chain output {
                type filter hook output priority 0;
        }
}'

# Path to the nftables configuration file based on the distribution
if [[ "$DISTRO" == "CentOS" || "$DISTRO" == "Red" ]]; then
    NFTABLES_CONF="/etc/sysconfig/nftables.conf"
else
    NFTABLES_CONF="/etc/nftables.conf"
fi

# Check if the required configuration is present in the nftables configuration file
if ! grep -q "flush ruleset" "$NFTABLES_CONF"; then
    # If not, add the required configuration to the nftables configuration file
    echo "$REQUIRED_CONF" > "$NFTABLES_CONF"
fi

# Check if nftables is enabled and start it if necessary
if ! systemctl is-enabled nftables &>/dev/null; then
    echo "Enabling nftables service..."
    systemctl enable nftables
fi

if ! systemctl is-active nftables &>/dev/null; then
    echo "Starting nftables service..."
    systemctl start nftables
fi
# Manages the services to run nPerfServer
# Script for Unix/Linux. Version : 2.1
# https://www.nperf.com/

OPT_ARGS=""
BIND_IP=""
NFT_ARGS=""
IPV6=""
NB_THREAD=""
DAEMON_START=""
PORT="8080"
TLS_PORT="8443"
REDIRECT_PORT_80=0
REDIRECT_PORT_443=0
IPV4_BIND=""
IPV6_BIND=""
IPV4_ADDRESSES=""
IPV6_ADDRESSES=""
LOG_CONNECTIONS=0

KERNEL_VERSION_MAJ=$(uname -r | cut -d '-' -f 1 | cut -d '.' -f 1)
KERNEL_VERSION_MIN=$(uname -r | cut -d '-' -f 1 | cut -d '.' -f 2)
if [ "$KERNEL_VERSION_MAJ" -ge "4" ]; then
    IPV6_REDIR=true
elif [ "$KERNEL_VERSION_MAJ" -ge "3" ] && [ "$KERNEL_VERSION_MIN" -ge "9" ]; then
    IPV6_REDIR=true
else
    IPV6_REDIR=false
fi

set -e

. "/etc/nperf/nperf-server.conf"

test -x $DAEMON || exit 1
test -z $1 && arg="empty" || arg=$1

set -u
DESC="`basename $DAEMON` daemon"

DAEMON_START="$DAEMON_START -p $PORT -t $TLS_PORT --uuidfile=$UUID"
if [ "$LOG_CONNECTIONS" -eq 1 ]; then
    DAEMON_START="$DAEMON_START -a"
fi
if [ -n "$OPT_ARGS" ]; then
    DAEMON_START="$DAEMON_START $OPT_ARGS"
fi


if [[ $BIND_IP = *":"* ]]; then
    IPV6=true
fi

TESTIP="localhost"
if [ -n "$BIND_IP" ]; then
    DAEMON_START="$DAEMON_START -i $BIND_IP"
    for address in $BIND_IP
    do
        if [[ $address = *":"* ]] ; then
            if [ "$address" != "::" ] ; then
                IPV6_ADDRESSES="$IPV6_ADDRESSES,$address/128"
                if [ -z "$TESTIP" ]; then
                    TESTIP=$address
                fi
            fi
        elif [ "$address" != "0.0.0.0" ]; then
            IPV4_ADDRESSES="$IPV4_ADDRESSES,$address/32"
            if [ -z "$TESTIP" ]; then
                TESTIP=$address
            fi
        fi
    done
    if [ -n "$IPV4_ADDRESSES" ] ; then
        NFT_ARGS=" ip daddr { ${IPV4_ADDRESSES:1} }"
        IPV4_BIND=" (binded to ${IPV4_ADDRESSES:1})"
    fi
    if [ -n "$IPV6_ADDRESSES" ] ; then
        NFT_ARGS=" ip6 daddr { ${IPV6_ADDRESSES:1} }"
        IPV6_BIND=" (binded to ${IPV6_ADDRESSES:1})"
    fi
fi

isrunning() {
        curl -s -m 5 -i $TESTIP:$PORT | grep "Server: nPerf" >/dev/null 2>&1 && return 0
        return 1
}

initnet() {
    echo -n "Activating Anti-DOS limitation..."
    CHECK_DOS=$(nft list ruleset | grep nPerfServer | grep "limit rate over" | cat)
    if [ -z "$CHECK_DOS" ]; then
        nft add rule inet filter input tcp flags syn limit rate over 260/second drop comment "nPerfServer"
        echo " OK."
    else
        echo " Already set!"
    fi

    echo "Flushing existing nPerfServer prerouting rules..."
    RULES=$(nft -a list ruleset)

    # Set a flag to track if we are within the inet nat table
    ip_nat_table=false
    ip6_nat_table=false

    # Iterate through each line of output
    while IFS= read -r line; do
        # Check if the line indicates the start of the inet nat table
        if [[ $line == *"table ip nat "* ]]; then
            ip_nat_table=true
        fi
        # Check if the line indicates the start of the ip6 nat table
        if [[ $line == *"table ip6 nat "* ]]; then
            ip6_nat_table=true
        fi

        # If we are within the inet nat table and the line contains "comment "nPerfServer""
        if [ "$ip_nat_table" = true ] && [[ $line == *"comment \"nPerfServer\""* ]]; then
            # Extract the handle number from the line using awk
            handle=$(echo "$line" | awk '{print $NF}')

            if [ -n "$handle" ] && [[ "$handle" =~ ^[0-9]+$ ]]; then
                echo "... Deleting rule with handle $handle"
                # Delete the rule using nft delete command
                nft delete rule ip nat prerouting handle "$handle"
            else
                echo "... Handle not found or invalid in line: $line"
            fi
        fi
        # If we are within the ip6 nat table and the line contains "comment "nPerfServer""
        if [ "$ip6_nat_table" = true ] && [[ $line == *"comment \"nPerfServer\""* ]]; then
            # Extract the handle number from the line using awk
            handle=$(echo "$line" | awk '{print $NF}')
            echo $handle

            if [ -n "$handle" ] && [[ "$handle" =~ ^[0-9]+$ ]]; then
                echo "... Deleting rule with handle $handle from ip6 nat table"
                # Delete the rule using nft delete command
                nft delete rule ip6 nat prerouting handle "$handle"
            else
                echo "... Handle not found or invalid in line: $line"
            fi
        fi
        # Check if the line indicates the end of the inet nat table
        if [[ $line == "}" && "$ip_nat_table" = true ]]; then
            ip_nat_table=false
        fi
        # Check if the line indicates the end of the ip6 nat table
        if [[ $line == "}" && "$ip6_nat_table" = true ]]; then
            ip6_nat_table=false
        fi
    done <<< "$RULES"
    # Create inet table and nat chain if not exists
    if ! nft list tables | grep -q 'ip  nat'; then
        nft add table ip nat
    fi
    if ! nft list tables | grep -q 'ip6  nat'; then
        nft add table ip6 nat
    fi
    if ! nft list table ip nat | grep -q 'chain prerouting'; then
        nft add chain ip nat prerouting '{type nat hook prerouting priority -101 ;}'
    fi
    if ! nft list table ip6 nat | grep -q 'chain prerouting'; then
        nft add chain ip6 nat prerouting '{type nat hook prerouting priority -101 ;}'
    fi
    sleep 1
    if [ "$REDIRECT_PORT_80" -eq 1 ]; then
        echo -n "Activating redirection from TCP port 80$IPV4_BIND$IPV6_BIND to port $PORT..."
        nft add rule ip nat prerouting tcp dport 80 redirect to $PORT comment "nPerfServer"
        nft add rule ip6 nat prerouting tcp dport 80 redirect to $PORT comment "nPerfServer"
        echo " OK."
    fi

    if [ "$REDIRECT_PORT_443" -eq 1 ]; then
        echo -n "Activating redirection from TCP port 443$IPV4_BIND$IPV6_BIND to port $TLS_PORT..."
        nft add rule ip nat prerouting tcp dport 443 redirect to $TLS_PORT comment "nPerfServer"
        nft add rule ip6 nat prerouting tcp dport 443 redirect to $TLS_PORT comment "nPerfServer"
        echo " OK."
    fi
}

start() {
    if [ -z "$NB_THREAD" ]; then
        echo -n "Starting $DESC ... "
    else
        echo -n "Starting $DESC with $NB_THREAD threads ... "
        DAEMON_START="$DAEMON_START -w $NB_THREAD"
    fi
    DAEMON_START="-x --pidfile=$PID $DAEMON_START"
    rm -f $PID
    set +e
    [ "`id -u`" -eq "0" ] && su $USER -s /bin/sh -c "${DAEMON} ${DAEMON_START}" >/dev/null
    [ "`id -u`" -eq "`id -u $USER`" ] && eval "${DAEMON} ${DAEMON_START}"
    set -e
    sleep 0.5 ; test -e $PID && echo "[ ok ]" || echo "[fail]"
}

stop() {
    echo -n "Stopping $DESC : "
    pkill `basename $DAEMON` || true
    sleep 0.5 ; test -e $PID && echo "[fail]" || echo "[ ok ]"
    echo -n "Flushing existing nPerfServer prerouting rules..."
    RULES=$(nft -a list ruleset)

    # Set a flag to track if we are within the inet nat table
    ip_nat_table=false
    ip6_nat_table=false

    # Iterate through each line of output
    while IFS= read -r line; do
        # Check if the line indicates the start of the inet nat table
        if [[ $line == *"table ip nat "* ]]; then
            ip_nat_table=true
        fi
        # Check if the line indicates the start of the ip6 nat table
        if [[ $line == *"table ip6 nat "* ]]; then
            ip6_nat_table=true
        fi

        # If we are within the inet nat table and the line contains "comment "nPerfServer""
        if [ "$ip_nat_table" = true ] && [[ $line == *"comment \"nPerfServer\""* ]]; then
            # Extract the handle number from the line using awk
            handle=$(echo "$line" | awk '{print $NF}')

            if [ -n "$handle" ] && [[ "$handle" =~ ^[0-9]+$ ]]; then
                echo "... Deleting rule with handle $handle"
                # Delete the rule using nft delete command
                nft delete rule ip nat prerouting handle "$handle"
            else
                echo "... Handle not found or invalid in line: $line"
            fi
        fi
        # If we are within the ip6 nat table and the line contains "comment "nPerfServer""
        if [ "$ip6_nat_table" = true ] && [[ $line == *"comment \"nPerfServer\""* ]]; then
            # Extract the handle number from the line using awk
            handle=$(echo "$line" | awk '{print $NF}')

            if [ -n "$handle" ] && [[ "$handle" =~ ^[0-9]+$ ]]; then
                echo "... Deleting rule with handle $handle from ip6 nat table"
                # Delete the rule using nft delete command
                nft delete rule ip6 nat prerouting handle "$handle"
            else
                echo "... Handle not found or invalid in line: $line"
            fi
        fi
        # Check if the line indicates the end of the inet nat table
        if [[ $line == "}" && "$ip_nat_table" = true ]]; then
            ip_nat_table=false
        fi
        # Check if the line indicates the end of the ip6 nat table
        if [[ $line == "}" && "$ip6_nat_table" = true ]]; then
            ip6_nat_table=false
        fi
    done <<< "$RULES"
}
case "$arg" in
    initnet)
        initnet
    ;;
    startcmd)
        if [ -n "$NB_THREAD" ]; then
            DAEMON_START="$DAEMON_START -w $NB_THREAD"
        fi
        echo "$DAEMON $DAEMON_START"
    ;;
    start-systemd)
        if [ -z "$NB_THREAD" ]; then
            echo -n "Starting $DESC ... "
        else
            echo -n "Starting $DESC with $NB_THREAD threads ... "
            DAEMON_START="$DAEMON_START -w $NB_THREAD"
        fi
        DAEMON_START="-x --pidfile=$PID $DAEMON_START"
        rm -f $PID
        set +e
        ${DAEMON} ${DAEMON_START}
        set -e
        exit $?
    ;;
    start)
        if isrunning; then
            echo "$DESC is already running. Use \"restart\" if you want to restart it."
        else
            initnet
            start
        fi
    ;;
    stop)
        if isrunning; then
            stop
        else
            echo "$DESC is not running. No need to stop it."
        fi
    ;;
    status)
        isrunning && echo "$DESC is running." || echo "$DESC is stopped."
    ;;
    version)
        $DAEMON -v || true
    ;;
    restart)
        if isrunning; then
            stop
            initnet
            start
        else
            echo "$DESC is not running. Use \"start\" if you want to start it."
        fi
    ;;
    setup-initd-autostart)
        echo -n "Setting up init.d to auto-start nPerfServer : "
        chkconfig --level 345 nperf-server on
        cat <<EOF > /etc/cron.d/nperf-server
# Persistent run (auto restart nPerfServer on crash)
*/2 * * * *    root   /etc/init.d/nperf-server status | grep "nPerfServer daemon is running" >/dev/null || /etc/init.d/nperf-server start
EOF
        echo "[ ok ]"
    ;;
    *)
        echo "Usage: $0 {start|stop|restart|status|version|setup-initd-autostart}" >&2
        exit 1
    ;;
esac

exit 0