#!/usr/bin/env bash
# ==========================================================================
#         _   _      _   ____            __ __  __      _
#        | \ | | ___| |_|  _ \ ___ _ __ / _|  \/  | ___| |_ ___ _ __
#        |  \| |/ _ \ __| |_) / _ \ '__| |_| |\/| |/ _ \ __/ _ \ '__|
#        | |\  |  __/ |_|  __/  __/ |  |  _| |  | |  __/ ||  __/ |
#        |_| \_|\___|\__|_|   \___|_|  |_| |_|  |_|\___|\__\___|_|
#
#                  NetPerfMeter -- Network Performance Meter
#                 Copyright (C) 2009-2026 by Thomas Dreibholz
# ==========================================================================
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
# Contact:  dreibh@simula.no
# Homepage: https://www.nntb.no/~dreibh/netperfmeter/

# Bash options:
set -euo pipefail


# ###### Load modules #######################################################
tryToLoadModules()
{
   local component="$1"
   local pattern="$2"
   shift 2

   echo -en "\x1b[32mLoading ${component} modules ... "
   while [ $# -gt 0 ] ; do
      local directory="$1"
      shift
      if [ -d "${directory}" ] ; then
         modules="$(find "${directory}" -name "${pattern}" -print0 | xargs -0 -r -n1 basename | sed -e "s/\.ko.*$//g" | sort)"
         for module in ${modules} ; do
            echo -en "\x1b[37m${module}\x1b[0m "
            if [ "${UNAME}" == "Linux" ] ; then
               modprobe "${module}" >/dev/null 2>&1
            elif [ "${UNAME}" == "FreeBSD" ] ; then
               kldload "${module}" >/dev/null 2>&1 || true
            elif [ "${UNAME}" == "NetBSD" ] ; then
               modload "${module}" >/dev/null 2>&1 || true
            fi
         done
      fi
   done
   echo -e "\x1b[0m"
}


# ###### Unload modules #####################################################
tryToUnloadModules()
{
   local component="$1"
   local pattern="$2"
   shift 2

   while [ $# -gt 0 ] ; do
      local directory="$1"
      shift
      if [ -d "${directory}" ] ; then
         modules="$(find "${directory}" -name "${pattern}" -print0 | xargs -0 -r -n1 basename | sed -e "s/\.ko.*$//g" | sort)"
         for module in ${modules} ; do
            if [ "${UNAME}" == "Linux" ] ; then
               rmmod "${module}" >/dev/null 2>&1 || true
            elif [ "${UNAME}" == "FreeBSD" ] ; then
               kldunload "${module}" >/dev/null 2>&1 || true
            elif [ "${UNAME}" == "NetBSD" ] ; then
               modunload "${module}" >/dev/null 2>&1 || true
            fi
         done
      fi
   done
}


# ###### Print status #######################################################
showStatus()
{
   local modules=""
   local allowedCCs=""
   local defaultCC=""

   # ------ Print the relevant modules --------------------------------------
   if [ "${UNAME}" == "Linux" ] ;  then
      modules="$(lsmod | ( grep "^tcp_\|^mptcp_\|^sctp\|^dccp|^quic" || true ) | awk '{ print $1 }' | sort | xargs)"
      allowedCCs="$(sysctl -n net.ipv4.tcp_allowed_congestion_control)"
      defaultCC="$(sysctl -n net.ipv4.tcp_congestion_control)"
   elif [ "${UNAME}" == "FreeBSD" ] ;  then
      modules="$(kldstat | ( grep -E "cc_" || true ) | awk '{ print $5 }' | sort | xargs)"
      allowedCCs="$(sysctl -n net.inet.tcp.cc.available | awk '(($1 != "CCmod") && ($1 != "")) { print $1 }' | xargs)"
      defaultCC="$(sysctl -n net.inet.tcp.cc.algorithm)"
   elif [ "${UNAME}" == "NetBSD" ] ; then
      modules="$(modstat | ( grep -E "sctp" || true ) | awk '{ print $1 }' | sort | xargs)"
      allowedCCs="$(sysctl -n net.inet.tcp.congctl.available)"
      defaultCC="$(sysctl -n net.inet.tcp.congctl.selected)"
   elif [ "${UNAME}" == "OpenBSD" ] ;  then
      modules="None (monolithic kernel)"
      allowedCCs="newreno"
      defaultCC="newreno"
   fi

   local numberOfAllowedCCs
   numberOfAllowedCCs="$(echo "${allowedCCs}" | xargs -n1 | wc -l | tr -d " ")"

   echo -en "\x1b[34m"
   echo -e "Modules:           \x1b[35m${modules}\x1b[34m"
   echo -e "Available TCP CCs: \x1b[35m${allowedCCs} \x1b[33m(total: ${numberOfAllowedCCs})\x1b[34m"
   echo -e "Default TCP CC:    \x1b[35m${defaultCC}\x1b[34m"
   echo -en "\x1b[0m"
}



# ###### Main program #######################################################

# ====== Check argument =====================================================
if [ $# -lt 1 ] ; then
  echo >&2 "Usage: $0 start|stop|status"
  exit 1
fi
UNAME="$(uname)"

# ====== Load modules =======================================================
if [ "$1" == "start" ] ; then

   if [ "${UNAME}" == "Linux" ] ; then
      kernel="$(uname -r)"
      tryToLoadModules "TCP"   "tcp_*.ko*"   "/lib/modules/${kernel}/kernel/net/ipv4/"  "/lib/modules/${kernel}/extra/net/ipv4/"
      tryToLoadModules "MPTCP" "mptcp_*.ko*" "/lib/modules/${kernel}/kernel/net/mptcp/" "/lib/modules/${kernel}/extra/net/mptcp/"
      tryToLoadModules "SCTP"  "sctp*.ko*"   "/lib/modules/${kernel}/kernel/net/sctp/"  "/lib/modules/${kernel}/extra/net/sctp/"
      tryToLoadModules "DCCP"  "dccp*.ko*"   "/lib/modules/${kernel}/kernel/net/dccp/"  "/lib/modules/${kernel}/extra/net/dccp/"
      tryToLoadModules "QUIC"  "quic.ko*"    "/lib/modules/${kernel}/kernel/net/quic/"  "/lib/modules/${kernel}/extra/"

      # Allow TCP CCs:
      availableCCs="$(sysctl net.ipv4.tcp_available_congestion_control | sed -e "s/.*= //g")"
      sysctl -qw net.ipv4.tcp_allowed_congestion_control="${availableCCs}"

   elif [ "${UNAME}" == "FreeBSD" ] ; then
      tryToLoadModules "CC"   "cc_*.ko" "/boot/kernel/"
      tryToLoadModules "SCTP" "sctp.ko" "/boot/kernel/"

  elif [ "${UNAME}" == "NetBSD" ] ; then
      true
      # arch="$(uname -m)"
      # kernel="$(uname -r)"
      # tryToLoadModules "SCTP" "sctp.kmod" "/stand/${arch}/${kernel}/modules/"

   elif [ "${UNAME}" == "OpenBSD" ] ; then
      true

   fi

   showStatus

# ====== Unload modules =====================================================
elif [ "$1" == "stop" ] ; then

   kernel="$(uname -r)"

   # Unload protocol modules, in two stages to cover dependencies:
   # shellcheck disable=SC2034
   for stage in 1 2 ; do
      if [ "${UNAME}" == "Linux" ] ; then
         tryToUnloadModules "TCP"   "tcp_*.ko*"   "/lib/modules/${kernel}/kernel/net/ipv4/"  "/lib/modules/${kernel}/extra/net/ipv4/"
         tryToUnloadModules "MPTCP" "mptcp_*.ko*" "/lib/modules/${kernel}/kernel/net/mptcp/" "/lib/modules/${kernel}/extra/net/mptcp/"
         tryToUnloadModules "SCTP"  "sctp*.ko*"   "/lib/modules/${kernel}/kernel/net/sctp/"  "/lib/modules/${kernel}/extra/net/sctp/"
         tryToUnloadModules "DCCP"  "dccp*.ko*"   "/lib/modules/${kernel}/kernel/net/dccp/"  "/lib/modules/${kernel}/extra/net/dccp/"
         tryToUnloadModules "QUIC"  "quic.ko*"    "/lib/modules/${kernel}/kernel/net/quic/"  "/lib/modules/${kernel}/extra/"

      elif [ "${UNAME}" == "FreeBSD" ] ; then
         tryToUnloadModules "CC"   "cc_*.ko" "/boot/kernel/"
         tryToUnloadModules "SCTP" "sctp.ko" "/boot/kernel/"

      elif [ "${UNAME}" == "NetBSD" ] ; then
         true
         # arch="$(uname -m)"
         # tryToUnloadModules "SCTP" "sctp.kmod" "/stand/${arch}/${kernel}/modules/"

      elif [ "${UNAME}" == "OpenBSD" ] ; then
         true
      fi
   done

   showStatus

# ====== Show status ========================================================
elif [ "$1" == "status" ] ; then

   showStatus

# ====== Invalid argument ===================================================
else
   echo >&2 "ERROR: Unknown command $1!"
   exit 1
fi
