#!/bin/bash

######################################################################################
#### Version 2.00                                                                 ####
#### For questions or comments contact@mylinux.work                               ####
#### Author : Phil Connor                                                         ####
####                                                                              ####
#### Notes :                                                                      ####
####   This script configures sysctl network and system tuning on Linux servers.  ####
####   It writes an idempotent drop-in file to /etc/sysctl.d/ so it can be run    ####
####   multiple times safely without duplicating settings.                         ####
####                                                                              ####
####   Use this script at your OWN risk. There is no guarantee whatsoever.        ####
####                                                                              ####
####   Usage:                                                                     ####
####     networktuning.sh                  - Standard tuning (HDD)                ####
####     networktuning.sh ssd              - SSD-optimized dirty page settings    ####
####     networktuning.sh nvme             - NVMe-optimized dirty page settings   ####
####     networktuning.sh --bbr            - Enable BBR congestion control        ####
####     networktuning.sh ssd --bbr        - SSD + BBR                            ####
####     networktuning.sh --router         - Enable IP forwarding                 ####
####     networktuning.sh --dry-run        - Show config without applying         ####
####     networktuning.sh --rollback       - Restore previous configuration       ####
######################################################################################

set -euo pipefail

##########################
#### System Variables ####
##########################
HOST=$(hostname)
CONF_FILE="/etc/sysctl.d/99-network-tuning.conf"
BACKUP_DIR="/etc/sysctl.d/backups"
STORAGE_TYPE="hdd"
ENABLE_BBR=false
ENABLE_ROUTER=false
DRY_RUN=false
ROLLBACK=false

#########################
#### Color Variables ####
#########################
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m'

######################
#### Root Check   ####
######################
if [[ $EUID -ne 0 ]]; then
    echo -e "${RED}Error: This script must be run as root${NC}"
    exit 1
fi

##############################
#### Parse CLI Arguments  ####
##############################
for arg in "$@"; do
    case "$arg" in
        ssd)       STORAGE_TYPE="ssd" ;;
        nvme)      STORAGE_TYPE="nvme" ;;
        --bbr)     ENABLE_BBR=true ;;
        --router)  ENABLE_ROUTER=true ;;
        --dry-run) DRY_RUN=true ;;
        --rollback) ROLLBACK=true ;;
        --help|-h)
            echo "Usage: $(basename "$0") [ssd|nvme] [--bbr] [--router] [--dry-run] [--rollback]"
            echo ""
            echo "Options:"
            echo "  ssd          Optimize dirty page settings for SSD storage"
            echo "  nvme         Optimize dirty page settings for NVMe storage"
            echo "  --bbr        Enable BBR congestion control (requires kernel 4.9+)"
            echo "  --router     Enable IP forwarding (IPv4 and IPv6)"
            echo "  --dry-run    Display the configuration without applying it"
            echo "  --rollback   Restore the previous configuration backup"
            echo "  --help       Show this help message"
            exit 0
            ;;
        *)
            echo -e "${RED}Unknown option: $arg${NC}"
            echo "Use --help for usage information"
            exit 1
            ;;
    esac
done

##################################
#### Detect OS and OS Version ####
##################################
if [[ ! -f /etc/os-release ]]; then
    echo -e "${RED}Error: /etc/os-release not found. Unsupported system.${NC}"
    exit 1
fi

OS_ID=$(. /etc/os-release && echo "${ID}")
OS_VERSION=$(. /etc/os-release && echo "${VERSION_ID}" | cut -d. -f1)

# shellcheck disable=SC2034
case "$OS_ID" in
    ubuntu|debian)
        PAKMGR="apt-get -y"
        ;;
    centos|rhel|oracle|rocky|almalinux|fedora)
        if [[ "$OS_VERSION" -le 7 ]]; then
            PAKMGR="yum -y"
        else
            PAKMGR="dnf -y"
        fi
        ;;
    *)
        echo -e "${YELLOW}Warning: Unrecognized OS '${OS_ID}'. Proceeding without package manager.${NC}"
        PAKMGR=""
        ;;
esac

################################
#### Handle Rollback        ####
################################
if [[ "$ROLLBACK" == true ]]; then
    LATEST_BACKUP=$(find "$BACKUP_DIR" -name "99-network-tuning.conf.*" -type f 2>/dev/null | sort -r | head -1)
    if [[ -z "$LATEST_BACKUP" ]]; then
        echo -e "${RED}Error: No backup found in $BACKUP_DIR${NC}"
        exit 1
    fi
    echo -e "${GREEN}Restoring from: $LATEST_BACKUP${NC}"
    cp "$LATEST_BACKUP" "$CONF_FILE"
    sysctl --system > /dev/null 2>&1
    echo -e "${GREEN}Rollback complete. Settings applied.${NC}"
    exit 0
fi

##########################
#### Sysctl Variables ####
##########################
MEM_BYTES=$(awk '/MemTotal:/ { printf "%0.f",$2 * 1024}' /proc/meminfo)
MEM_KB=$(awk '/MemTotal:/ { printf "%0.f",$2}' /proc/meminfo)
MAX_ORPHAN=$((MEM_BYTES / 10 / 65536))
FILE_MAX=$((MEM_BYTES / 4194304 * 256))
MAX_TW=$((FILE_MAX * 2))
MIN_FREE=$((MEM_KB / 100))

######################################
#### Storage-based dirty settings ####
######################################
case "$STORAGE_TYPE" in
    hdd)
        VM_DIRTY_BG_RATIO=5
        VM_DIRTY_RATIO=15
        ;;
    ssd)
        VM_DIRTY_BG_RATIO=3
        VM_DIRTY_RATIO=5
        ;;
    nvme)
        VM_DIRTY_BG_RATIO=1
        VM_DIRTY_RATIO=3
        ;;
esac

####################################
#### BBR availability check     ####
####################################
if [[ "$ENABLE_BBR" == true ]]; then
    if ! modprobe tcp_bbr 2>/dev/null; then
        echo -e "${YELLOW}Warning: BBR kernel module not available. Skipping BBR configuration.${NC}"
        ENABLE_BBR=false
    fi
fi

#################################
#### Generate Configuration  ####
#################################
echo "#######################################"
echo "#### Network Tuning for $HOST"
echo "#### Storage: $STORAGE_TYPE | BBR: $ENABLE_BBR | Router: $ENABLE_ROUTER"
echo "#######################################"

CONFIG=$(cat << EOF
###############################################################################
# Network Tuning Configuration
# Generated: $(date '+%Y-%m-%d %H:%M:%S')
# Host: $HOST
# Storage: $STORAGE_TYPE | BBR: $ENABLE_BBR | Router: $ENABLE_ROUTER
# Generator: networktuning.sh v2.00
###############################################################################

############################
#### Performance Tuning ####
############################

# Enable SYN cookies for SYN flood protection
net.ipv4.tcp_syncookies = 1

# Basic TCP tuning
net.ipv4.tcp_keepalive_time = 600
net.ipv4.tcp_synack_retries = 3
net.ipv4.tcp_syn_retries = 3

# RFC1337 — handle TIME-WAIT assassination hazards
net.ipv4.tcp_rfc1337 = 1

# Defines the local port range used by TCP and UDP
net.ipv4.ip_local_port_range = 1024 65535

# Log Martian packets with impossible addresses
net.ipv4.conf.all.log_martians = 1
net.ipv4.conf.default.log_martians = 1

# Enable window scaling (RFC1323)
net.ipv4.tcp_window_scaling = 1

# Enable timestamps (RFC1323)
net.ipv4.tcp_timestamps = 1

# Enable selective acknowledgments
net.ipv4.tcp_sack = 1

# Allow TCP to send duplicate SACKs
net.ipv4.tcp_dsack = 1

# Loose reverse path filtering (compatible with cloud/multihomed)
net.ipv4.conf.default.rp_filter = 2
net.ipv4.conf.all.rp_filter = 2

# Max remembered connection requests in SYN backlog
net.ipv4.tcp_max_syn_backlog = 20000

# Max orphaned TCP sockets (calculated from system memory)
net.ipv4.tcp_max_orphans = $MAX_ORPHAN

# Retries before killing an orphaned TCP connection
net.ipv4.tcp_orphan_retries = 1

# Time to hold sockets in FIN-WAIT-2 state (seconds)
net.ipv4.tcp_fin_timeout = 20

# Maximum TIME-WAIT sockets held simultaneously
net.ipv4.tcp_max_tw_buckets = $MAX_TW

# Don't cache ssthresh from previous connections
net.ipv4.tcp_no_metrics_save = 1

# Enable receive buffer auto-tuning
net.ipv4.tcp_moderate_rcvbuf = 1

# TCP buffer auto-tuning limits (min default max bytes)
net.ipv4.tcp_rmem = 4096 87380 16777216
net.ipv4.tcp_wmem = 4096 65536 16777216

# Max socket buffer sizes
net.core.rmem_max = 16777216
net.core.wmem_max = 16777216

# Network device backlog queue length
net.core.netdev_max_backlog = 2500

# Max pending connections in listen queue
net.core.somaxconn = 65000

############################
#### VM / Storage Tuning ###
############################

# Prefer keeping processes in memory over swapping (1 = minimal swap)
vm.swappiness = 1

# Dirty page writeback ratios (tuned for $STORAGE_TYPE)
# Monitor with: grep -A 1 dirty /proc/vmstat
vm.dirty_background_ratio = $VM_DIRTY_BG_RATIO
vm.dirty_ratio = $VM_DIRTY_RATIO

# Required free memory (1% of physical RAM)
vm.min_free_kbytes = $MIN_FREE

# System open file limit (calculated from system memory)
fs.file-max = $FILE_MAX

############################
#### Kernel Settings    ####
############################

# Kernel log levels
kernel.printk = 4 4 1 7
kernel.core_uses_pid = 1
kernel.sysrq = 0

# Disable core dumps for setuid programs
fs.suid_dumpable = 0

###########################
#### Security Settings ####
###########################

# Full ASLR (randomize stack, VDSO, mmap, heap)
kernel.randomize_va_space = 2

# Don't accept ICMP redirects
net.ipv4.conf.all.accept_redirects = 0
net.ipv4.conf.default.accept_redirects = 0
net.ipv6.conf.all.accept_redirects = 0
net.ipv6.conf.default.accept_redirects = 0

# Don't send ICMP redirects
net.ipv4.conf.all.send_redirects = 0
net.ipv4.conf.default.send_redirects = 0

# Don't accept IP source route packets
net.ipv4.conf.all.accept_source_route = 0
net.ipv4.conf.default.accept_source_route = 0
net.ipv6.conf.all.accept_source_route = 0

# Ignore ICMP broadcasts and bogus error responses
net.ipv4.icmp_echo_ignore_broadcasts = 1
net.ipv4.icmp_ignore_bogus_error_responses = 1

# IPv6 router advertisement hardening
net.ipv6.conf.default.accept_ra_rtr_pref = 0
net.ipv6.conf.all.accept_ra_rtr_pref = 0
net.ipv6.conf.default.accept_ra_pinfo = 0
net.ipv6.conf.all.accept_ra_pinfo = 0
net.ipv6.conf.default.accept_ra_defrtr = 0
net.ipv6.conf.all.accept_ra_defrtr = 0
EOF
)

# Append BBR congestion control if enabled
if [[ "$ENABLE_BBR" == true ]]; then
    CONFIG+=$(cat << 'EOF'


###########################
#### BBR Congestion    ####
###########################

# Use fair queuing packet scheduler (required for BBR)
net.core.default_qdisc = fq

# Enable BBR congestion control algorithm
net.ipv4.tcp_congestion_control = bbr
EOF
)
fi

# Append router/forwarding settings if enabled
if [[ "$ENABLE_ROUTER" == true ]]; then
    CONFIG+=$(cat << 'EOF'


###########################
#### Router / Forward  ####
###########################

# Enable IP forwarding (IPv4 and IPv6)
net.ipv4.ip_forward = 1
net.ipv6.conf.all.forwarding = 1
EOF
)
fi

##############################
#### Apply Configuration  ####
##############################
if [[ "$DRY_RUN" == true ]]; then
    echo ""
    echo -e "${YELLOW}--- Dry Run: Configuration Preview ---${NC}"
    echo "$CONFIG"
    echo -e "${YELLOW}--- End Preview ---${NC}"
    exit 0
fi

# Create backup directory
mkdir -p "$BACKUP_DIR"

# Backup existing config if present
if [[ -f "$CONF_FILE" ]]; then
    cp "$CONF_FILE" "${BACKUP_DIR}/99-network-tuning.conf.$(date +%Y%m%d-%H%M%S)"
    echo -e "${GREEN}Backed up existing config to $BACKUP_DIR${NC}"
fi

# Write configuration atomically
TEMP_FILE=$(mktemp)
echo "$CONFIG" > "$TEMP_FILE"

# Validate before installing
if sysctl -p "$TEMP_FILE" > /dev/null 2>&1; then
    mv "$TEMP_FILE" "$CONF_FILE"
    chmod 644 "$CONF_FILE"
    sysctl --system > /dev/null 2>&1
    echo -e "${GREEN}Network tuning applied successfully to $CONF_FILE${NC}"
else
    echo -e "${RED}Error: Configuration validation failed. No changes applied.${NC}"
    echo -e "${RED}Review errors above and check kernel compatibility.${NC}"
    rm -f "$TEMP_FILE"
    exit 1
fi

echo -e "${GREEN}Done. Use 'sysctl --system' to verify or '$(basename "$0") --rollback' to revert.${NC}"
exit 0