#!/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 || true
            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 -E "^(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}" | wc -w | 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"
}


# ###### Version ############################################################
version () {
   echo "netperfmeter-module-loader @BUILD_MAJOR@.@BUILD_MINOR@.@BUILD_PATCH@"
   exit 0
}


# ------ BEGIN-OF-USAGE-BLOCK ------
# ###### Usage ##############################################################
usage () {
   local program="$1"
   local exitCode="$2"
   local header=$'\x1b[4m'
   local cmd=$'\x1b[1m'
   local opt=$'\x1b[1m'
   # local arg=$'\x1b[4m'
   local syn=$'\x1b[1m'
   local reset=$'\x1b[0m'
   cat >&2 <<EOF
${header}Usage:${reset}
${cmd}${program}${reset} ${syn}start${reset}
${cmd}${program}${reset} ${syn}stop${reset}
${cmd}${program}${reset} ${syn}status${reset}
${cmd}${program}${reset}
	${syn}[${reset}${opt}-h${reset} ${syn}|${reset} ${opt}--help${reset}${syn}]${reset}
${cmd}${program}${reset}
	${syn}[${reset}${opt}-v${reset} ${syn}|${reset} ${opt}--version${reset}${syn}]${reset}
EOF
   exit "${exitCode}"
}
# ------ END-OF-USAGE-BLOCK ------



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

# ====== Handle arguments ===================================================
UNAME="$(uname)"
case "${UNAME}" in
   FreeBSD)
      GETOPT="$(PATH="/usr/local/bin:${PATH}" command -v getopt || true)"
      ;;
   NetBSD)
      GETOPT="$(PATH="/usr/pkg/bin:${PATH}" command -v getopt || true)"
      ;;
   OpenBSD)
      GETOPT="$(command -v gnugetopt || true)"
      ;;
   SunOS)
      GETOPT="$(PATH="/usr/gnu/bin:${PATH}" command -v getopt || true)"
      ;;
   Darwin)
      GETOPT="$(env PATH="/opt/homebrew/opt/gnu-getopt/bin:${PATH}" command -v getopt || true)"
      ;;
   *)
      GETOPT="$(command -v getopt || true)"
      ;;
esac
if [ -z "${GETOPT}" ] || "${GETOPT}" -T >/dev/null 2>&1 || [[ $? -ne 4 ]]; then
   echo >&2 "ERROR: Enhanced/GNU getopt is required!"
   exit 1
fi

if ! options="$("${GETOPT}" -o hv --long help,version -a -- "$@")"; then
   usage "$0" 1
fi

eval set -- "${options}"
while [ $# -gt 0 ] ; do
   case "$1" in
      -h | --help)
         usage "$0" 0
         ;;
      -v | --version)
         version
         ;;
      --)
         shift
         break
         ;;
   esac
   shift
done

if [ $# -ne 1 ] ; then
   usage "$0" 1
fi
COMMAND="$1"

# ====== Load modules =======================================================
if [ "${COMMAND}" == "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 -n net.ipv4.tcp_available_congestion_control)"
      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 [ "${COMMAND}" == "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 [ "${COMMAND}" == "status" ] ; then

   showStatus

# ====== Invalid argument ===================================================
else
   echo >&2 "ERROR: Unknown command ${COMMAND}!"
   usage "$0" 1
fi
