#!/usr/bin/env bash

# shellcheck disable=1117
# shellcheck disable=2162

# Enforce script only be run on Proxmox hosts
if [ ! -f "/etc/pve/.version" ] ; then
  echo "ERROR: This script only supports Proxmox"
  exit 1
fi

if [ -f "/etc/system-optimization" ] ; then
  echo "ERROR: Script can only be run once"
  exit 1
fi

function ask-user() {

# Check for AMD processors (EPYC or Ryzen)
echo -e "Check for AMD processors (EPYC or Ryzen)"
if [ "$(grep -i -m 1 "model name" /proc/cpuinfo | grep -i "EPYC")" != "" ]; then
    echo "AMD EPYC detected"
    if [ -z "$PMX_AMDFIXES" ]; then
        read -p "Do you want to enable AMD fixes? (yes/no): " PMX_AMDFIXES
    fi
elif [ "$(grep -i -m 1 "model name" /proc/cpuinfo | grep -i "Ryzen")" != "" ]; then
    echo "AMD Ryzen detected"
    if [ -z "$PMX_AMDFIXES" ]; then
        read -p "Do you want to enable AMD fixes? (yes/no): " PMX_AMDFIXES
    fi
else
    sleep 1
    echo -e "Non-AMD processor detected... Assuming Intel processor."
fi

# Force APT to use IPv4
if [ -z "$PMX_APTIPV4" ]; then
    while true; do
        read -p "Do you want to force APT to use IPv4? (yes/no): " PMX_APTIPV4
        case "$PMX_APTIPV4" in
            yes|no) break ;;
            *) echo "Please answer 'yes' or 'no'." ;;
        esac
    done
fi

# Update Proxmox and install various system utils
if [ -z "$PMX_APTUPGRADE" ]; then
    while true; do
        read -p "Do you want to update Proxmox and install various system utilities? (yes/no): " PMX_APTUPGRADE
        case "$PMX_APTUPGRADE" in
            yes|no) break ;;
            *) echo "Please answer 'yes' or 'no'." ;;
        esac
    done
fi

# Customize bashrc
if [ -z "$PMX_BASHRC" ]; then
    while true; do
        read -p "Do you want to customize bashrc? (yes/no): " PMX_BASHRC
        case "$PMX_BASHRC" in
            yes|no) break ;;
            *) echo "Please answer 'yes' or 'no'." ;;
        esac
    done
fi

# Add the latest Ceph provided by Proxmox
if [ -z "$PMX_CEPH" ]; then
    while true; do
        read -p "Do you want to add the latest Ceph provided by Proxmox? (yes/no): " PMX_CEPH
        case "$PMX_CEPH" in
            yes|no) break ;;
            *) echo "Please answer 'yes' or 'no'." ;;
        esac
    done
fi

# Disable portmapper / rpcbind for security
if [ -z "$PMX_DISABLERPC" ]; then
    while true; do
        read -p "Do you want to disable portmapper/rpcbind for security? (yes/no): " PMX_DISABLERPC
        case "$PMX_DISABLERPC" in
            yes|no) break ;;
            *) echo "Please answer 'yes' or 'no'." ;;
        esac
    done
fi

# Ensure Entropy Pools are Populated
if [ -z "$PMX_ENTROPY" ]; then
    while true; do
        read -p "Do you want to ensure entropy pools are populated? (yes/no): " PMX_ENTROPY
        case "$PMX_ENTROPY" in
            yes|no) break ;;
            *) echo "Please answer 'yes' or 'no'." ;;
        esac
    done
fi

# Protect the web interface with fail2ban
if [ -z "$PMX_FAIL2BAN" ]; then
    while true; do
        read -p "Do you want to protect the web interface with fail2ban? (yes/no): " PMX_FAIL2BAN
        case "$PMX_FAIL2BAN" in
            yes|no) break ;;
            *) echo "Please answer 'yes' or 'no'." ;;
        esac
    done
fi

# Install ifupdown2 for a virtual internal network
if [ -z "$PMX_IFUPDOWN2" ]; then
    while true; do
        read -p "Do you want to install ifupdown2 for a virtual internal network? (yes/no): " PMX_IFUPDOWN2
        case "$PMX_IFUPDOWN2" in
            yes|no) break ;;
            *) echo "Please answer 'yes' or 'no'." ;;
        esac
    done
fi

# Limit the size and optimize journald
if [ -z "$PMX_JOURNALD" ]; then
    while true; do
        read -p "Do you want to limit the size and optimize journald? (yes/no): " PMX_JOURNALD
        case "$PMX_JOURNALD" in
            yes|no) break ;;
            *) echo "Please answer 'yes' or 'no'." ;;
        esac
    done
fi

# Install kernel source headers
if [ -z "$PMX_KERNELHEADERS" ]; then
    while true; do
        read -p "Do you want to install kernel source headers? (yes/no): " PMX_KERNELHEADERS
        case "$PMX_KERNELHEADERS" in
            yes|no) break ;;
            *) echo "Please answer 'yes' or 'no'." ;;
        esac
    done
fi

# Ensure ksmtuned is enabled and optimize according to RAM size
if [ -z "$PMX_KSMTUNED" ]; then
    while true; do
        read -p "Do you want to ensure ksmtuned is enabled and optimize according to RAM size? (yes/no): " PMX_KSMTUNED
        case "$PMX_KSMTUNED" in
            yes|no) break ;;
            *) echo "Please answer 'yes' or 'no'." ;;
        esac
    done
fi

# Set language to en_US.UTF-8
if [ -z "$PMX_LANG" ]; then
    while true; do
        read -p "Do you want to set the language to en_US.UTF-8? (yes/no): " PMX_LANG
        case "$PMX_LANG" in
            yes|no) break ;;
            *) echo "Please answer 'yes' or 'no'." ;;
        esac
    done
fi

# Enable restart on kernel panic, kernel oops, and hardlockup
if [ -z "$PMX_KERNELPANIC" ]; then
    while true; do
        read -p "Do you want to enable restart on kernel panic, kernel oops, and hardlockup? (yes/no): " PMX_KERNELPANIC
        case "$PMX_KERNELPANIC" in
            yes|no) break ;;
            *) echo "Please answer 'yes' or 'no'." ;;
        esac
    done
fi

# Increase max user watches, FD limit, FD ulimit, and max key limit
if [ -z "$PMX_LIMITS" ]; then
    while true; do
        read -p "Do you want to increase max user watches, FD limit, FD ulimit, and max key limit? (yes/no): " PMX_LIMITS
        case "$PMX_LIMITS" in
            yes|no) break ;;
            *) echo "Please answer 'yes' or 'no'." ;;
        esac
    done
fi

# Optimize logrotate
if [ -z "$PMX_LOGROTATE" ]; then
    while true; do
        read -p "Do you want to optimize logrotate? (yes/no): " PMX_LOGROTATE
        case "$PMX_LOGROTATE" in
            yes|no) break ;;
            *) echo "Please answer 'yes' or 'no'." ;;
        esac
    done
fi

# 2FA for SSH and TUI
if [ -z "$PMX_2FA" ]; then
    while true; do
        read -p "Do you want 2FA for SSH and TUI? (yes/no): " PMX_2FA
        case "$PMX_2FA" in
            yes|no) break ;;
            *) echo "Please answer 'yes' or 'no'." ;;
        esac
    done
fi

# Summary Section
echo -e "\nSummary of your answers:"
echo "PMX_AMDFIXES: ${PMX_AMDFIXES:-not set}"
echo "PMX_APTIPV4: ${PMX_APTIPV4:-not set}"
echo "PMX_APTUPGRADE: ${PMX_APTUPGRADE:-not set}"
echo "PMX_BASHRC: ${PMX_BASHRC:-not set}"
echo "PMX_CEPH: ${PMX_CEPH:-not set}"
echo "PMX_DISABLERPC: ${PMX_DISABLERPC:-not set}"
echo "PMX_ENTROPY: ${PMX_ENTROPY:-not set}"
echo "PMX_FAIL2BAN: ${PMX_FAIL2BAN:-not set}"
echo "PMX_IFUPDOWN2: ${PMX_IFUPDOWN2:-not set}"
echo "PMX_JOURNALD: ${PMX_JOURNALD:-not set}"
echo "PMX_KERNELHEADERS: ${PMX_KERNELHEADERS:-not set}"
echo "PMX_KSMTUNED: ${PMX_KSMTUNED:-not set}"
echo "PMX_LANG: ${PMX_LANG:-not set}"
echo "PMX_KERNELPANIC: ${PMX_KERNELPANIC:-not set}"
echo "PMX_LIMITS: ${PMX_LIMITS:-not set}"
echo "PMX_LOGROTATE: ${PMX_LOGROTATE:-not set}"
echo "PMX_2FA: ${PMX_2FA:-not set}"

}

## Main Process

function main() {
echo "Lets start with your chosen optimizations .... "

# Check if PMX_2FA is not set or is empty
if [ "${PMX_2FA,,}" == "yes" ] ; then
	echo "2FA for SSH & TUI are not set, proceeding with 2FA setup."
    # Update system packages
    apt update && sudo apt upgrade -y
    # Install necessary packages
    apt install -y libpam-google-authenticator ssh
    # Install Google Authenticator PAM module
    apt install -y libpam-google-authenticator
    # Configure SSH for 2FA
    # Backup the original sshd_config file
    cp /etc/ssh/sshd_config /etc/ssh/sshd_config.bak
    # Modify sshd_config to enable ChallengeResponseAuthentication
    sed -i 's/^#ChallengeResponseAuthentication no/ChallengeResponseAuthentication yes/' /etc/ssh/sshd_config
    # Restart the SSH service to apply changes
    systemctl restart sshd
    # Set up Google Authenticator for the user
    echo "Setting up Google Authenticator for user $(whoami)..."
    # Run the Google Authenticator setup
    google-authenticator
    # Enable PAM Google Authenticator for TUI login
    echo "Configuring TUI login with Google Authenticator..."
    # Backup PAM's common-auth configuration file
    cp /etc/pam.d/common-auth /etc/pam.d/common-auth.bak
    # Add Google Authenticator PAM module to common-auth
    echo "auth required pam_google_authenticator.so" | sudo tee -a /etc/pam.d/common-auth
    # Apply changes to PAM configuration
    systemctl restart systemd-logind
    # Output success message
    echo "2FA setup complete! Please verify by logging in via SSH or TUI."
    # Inform user of next steps
    echo "1. Use the Google Authenticator app to scan the QR code during the setup."
    echo "2. After completing the setup, try logging in via SSH or a TTY session."
else
    echo "PMX_2FA is set, skipping 2FA setup."
fi

# Set the local
if [ "$PMX_LANG" == "yes" ] ; then
    PMX_LANG="en_US.UTF-8"
fi
echo -r "Set the locales"
export LANG="$PMX_LANG"
export LC_ALL="C"

# Get variables
echo -e "Get current OS information"
OS_CODENAME="$(grep "VERSION_CODENAME=" /etc/os-release | cut -d"=" -f 2 | xargs )"
echo -e "Get current RAM configuration"
RAM_SIZE_GB=$(( $(vmstat -s | grep -i "total memory" | xargs | cut -d" " -f 1) / 1024 / 1000))

if [ "${PMX_LANG}" == "en_US.UTF-8" ] && [ "${PMX_NOAPTLANG,,}" == "yes" ] ; then
    # save bandwidth and skip downloading additional languages
    echo -e "Disable translations:"
    echo -e "Save bandwidth and skip downloading additional languages."
    echo -e "Acquire::Languages \"none\";\\n" > /etc/apt/apt.conf.d/99-pmx-disable-translations
fi

if [ "${PMX_APTIPV4,,}" == "yes" ] ; then
    # force APT to use IPv4
    echo -e "Force APT to use IPv4."
    echo -e "Acquire::ForceIPv4 \"true\";\\n" > /etc/apt/apt.conf.d/99-pmx-force-ipv4
fi

if [ "${PMX_NOENTREPO,,}" == "yes" ] ; then
    # Disable enterprise proxmox repo
    if [ -f /etc/apt/sources.list.d/pve-enterprise.list ]; then
      echo -e "Disable Enterprise Proxmox Repo."
      sed -i "s/^deb/#deb/g" /etc/apt/sources.list.d/pve-enterprise.list
    fi
    # Enable free public proxmox repo
    if [ ! -f /etc/apt/sources.list.d/proxmox.list ] && [ ! -f /etc/apt/sources.list.d/pve-public-repo.list ] && [ ! -f /etc/apt/sources.list.d/pve-install-repo.list ] ; then
      echo -e "Enable non-subscription Proxmox Repo."
      echo -e "deb http://download.proxmox.com/debian/pve ${OS_CODENAME} pve-no-subscription\\n" > /etc/apt/sources.list.d/pve-public-repo.list
    fi
    if [ "${PMX_TESTREPO,,}" == "yes" ] ; then
        # Enable Testing Proxmox Repo
        echo -e "Enable Testing Proxmox Repo."
        echo -e "deb http://download.proxmox.com/debian/pve ${OS_CODENAME} pvetest\\n" > /etc/apt/sources.list.d/pve-testing-repo.list
    fi
fi

# Rebuild and add non-free to /etc/apt/sources.list
echo -e "Rebuild and add non-free to /etc/apt/sources.list"
cat <<EOF > /etc/apt/sources.list
deb https://ftp.debian.org/debian ${OS_CODENAME} main contrib
deb https://ftp.debian.org/debian ${OS_CODENAME}-updates main contrib
# non-free
deb https://httpredir.debian.org/debian/ ${OS_CODENAME} main contrib non-free
# security updates
deb https://security.debian.org/debian-security ${OS_CODENAME}/updates main contrib
EOF

# Refresh the package lists
echo -e "Refresh the package lists..."
apt-get update > /dev/null 2>&1

# Remove conflicting utilities
echo -e "Refresh the package lists"
/usr/bin/env DEBIAN_FRONTEND=noninteractive apt-get -y -o Dpkg::Options::='--force-confdef' purge ntp openntpd systemd-timesyncd

# Fixes for common apt repo errors
echo -e "Refresh the package lists"
/usr/bin/env DEBIAN_FRONTEND=noninteractive apt-get -y -o Dpkg::Options::='--force-confdef' install apt-transport-https debian-archive-keyring ca-certificates curl

if [ "${PMX_APTUPGRADE,,}" == "yes" ] ; then
    # Update Proxmox and install a few imporant devops/sysadmin utils
    echo -e "Update Proxmox and install a few imporant devops/sysadmin utils"
    /usr/bin/env DEBIAN_FRONTEND=noninteractive apt-get -y -o Dpkg::Options::='--force-confdef' dist-upgrade
    pveam update
fi

# Install packages which are sometimes missing on Proxmox installs.
echo -e "Update Proxmox and install a few important devops/sysadmin utils"
/usr/bin/env DEBIAN_FRONTEND=noninteractive apt-get -y -o Dpkg::Options::='--force-confdef' install zfsutils-linux proxmox-backup-restore-image chrony

if [ "${PMX_UTILS,,}" == "yes" ] ; then
# Install common system utilities
    /usr/bin/env DEBIAN_FRONTEND=noninteractive apt-get -y -o Dpkg::Options::='--force-confdef' install \
    axel \
    build-essential \
    curl \
    dialog \
    dnsutils \
    dos2unix \
    git \
    gnupg-agent \
    grc \
    htop \
    iftop \
    iotop \
    iperf \
    ipset \
    iptraf \
    mlocate \
    msr-tools \
    nano \
    net-tools \
    omping \
    software-properties-common \
    sshpass \
    tmux \
    unzip \
    vim \
    vim-nox \
    wget \
    whois \
    zip \
    ncdu \
    duf
fi

if [ "${PMX_CEPH,,}" == "yes" ] ; then
    # Add the latest CEPH packages provided by Proxmox
    echo -e "Add the latest CEPH packages provided by Proxmox"
    echo "deb http://download.proxmox.com/debian/ceph-pacific ${OS_CODENAME} main" > /etc/apt/sources.list.d/ceph-pacific.list
    ## Refresh the package lists
    echo -e "Refresh the package lists..."
    apt-get update > /dev/null 2>&1
    ## Install CEPH support
    echo -e "Install CEPH support packages for Proxmox"
    echo "Y" | pveceph install
fi

if [ "${PMX_LYNIS,,}" == "yes" ] ; then
    # Lynis security scan tool by Cisofy
    echo -e "Installing Lynis security scan tool"
    wget -O - https://packages.cisofy.com/keys/cisofy-software-public.key | apt-key add -
    ## Add the latest lynis
    echo "deb https://packages.cisofy.com/community/lynis/deb/ stable main" > /etc/apt/sources.list.d/cisofy-lynis.list
    ## Refresh the package lists
    apt-get update > /dev/null 2>&1
    /usr/bin/env DEBIAN_FRONTEND=noninteractive apt-get -y -o Dpkg::Options::='--force-confdef' install lynis
fi

if [ "${PMX_OPENVSWITCH,,}" == "yes" ] && [ "${PMX_IFUPDOWN2}" == "no" ] ; then
    ## Install openvswitch for a virtual internal network
    echo -e "Install openvswitch for a virtual internal network"
    /usr/bin/env DEBIAN_FRONTEND=noninteractive apt-get -y -o Dpkg::Options::='--force-confdef' install ifenslave ifupdown
    /usr/bin/env DEBIAN_FRONTEND=noninteractive apt-get -y -o Dpkg::Options::='--force-confdef' remove ifupdown2
    /usr/bin/env DEBIAN_FRONTEND=noninteractive apt-get -y -o Dpkg::Options::='--force-confdef' install openvswitch-switch
else
    ## Install ifupdown2 for a virtual internal network allows rebootless networking changes (not compatible with openvswitch-switch)
    echo -e "Install ifupdown2 for a virtual internal network allows rebootless networking changes (not compatible with openvswitch-switch)"
    /usr/bin/env DEBIAN_FRONTEND=noninteractive apt-get -y -o Dpkg::Options::='--force-confdef' purge openvswitch-switch
    /usr/bin/env DEBIAN_FRONTEND=noninteractive apt-get -y -o Dpkg::Options::='--force-confdef' install ifupdown2
    /usr/bin/env DEBIAN_FRONTEND=noninteractive apt-get -y -o Dpkg::Options::='--force-confdef' remove ifenslave ifupdown
fi

if [ "${PMX_ZFSAUTOSNAPSHOT,,}" == "yes" ] ; then
    ## Install zfs-auto-snapshot
    echo -e "Install zfs-auto-snapshot utility"
    /usr/bin/env DEBIAN_FRONTEND=noninteractive apt-get -y -o Dpkg::Options::='--force-confdef' install zfs-auto-snapshot
    
    # Prompt user for snapshot retention settings
    read -p "How many 5-minute snapshots would you like to keep? " SNAP_5MIN
    read -p "How many hourly snapshots would you like to keep? " SNAP_HOURLY
    read -p "How many daily snapshots would you like to keep? " SNAP_DAILY
    read -p "How many weekly snapshots would you like to keep? " SNAP_WEEKLY
    read -p "How many monthly snapshots would you like to keep? " SNAP_MONTHLY

    # Make 5min snapshots, keep specified number
    if [ -f "/etc/cron.d/zfs-auto-snapshot" ] ; then
    sed -i "s|--keep=[0-9]*|--keep=$SNAP_5MIN|g" /etc/cron.d/zfs-auto-snapshot
    sed -i "s|*/[0-9]*|*/5|g" /etc/cron.d/zfs-auto-snapshot

    # Keep hourly snapshots, retain the specified number
    if [ -f "/etc/cron.hourly/zfs-auto-snapshot" ] ; then
      sed -i "s|--keep=[0-9]*|--keep=$SNAP_HOURLY|g" /etc/cron.hourly/zfs-auto-snapshot
    fi

    # Keep daily snapshots, retain the specified number
    if [ -f "/etc/cron.daily/zfs-auto-snapshot" ] ; then
      sed -i "s|--keep=[0-9]*|--keep=$SNAP_DAILY|g" /etc/cron.daily/zfs-auto-snapshot
    fi

    # Keep weekly snapshots, retain the specified number
    if [ -f "/etc/cron.weekly/zfs-auto-snapshot" ] ; then
      sed -i "s|--keep=[0-9]*|--keep=$SNAP_WEEKLY|g" /etc/cron.weekly/zfs-auto-snapshot
    fi

    # Keep monthly snapshots, retain the specified number
    if [ -f "/etc/cron.monthly/zfs-auto-snapshot" ] ; then
      sed -i "s|--keep=[0-9]*|--keep=$SNAP_MONTHLY|g" /etc/cron.monthly/zfs-auto-snapshot
    fi
  fi
fi

if [ "${PMX_KSMTUNED,,}" == "yes" ] ; then
    ## Ensure ksmtuned (ksm-control-daemon) is enabled and optimise according to ram size
    echo -e "Ensure ksmtuned (ksm-control-daemon) is enabled and optimise according to ram size"
    
    # Install ksm-control-daemon
    /usr/bin/env DEBIAN_FRONTEND=noninteractive apt-get -y -o Dpkg::Options::='--force-confdef' install ksm-control-daemon

    # Get the system's RAM size in GB
    RAM_SIZE_GB=$(free -g | awk '/^Mem:/{print $2}')

    # Ask user for KSM settings
    echo "Detected RAM size: $RAM_SIZE_GB GB"
    echo "Please specify the KSM settings for your system."

    # Get KSM threshold and sleep values from user
    read -p "Enter the KSM threshold coefficient (as a percentage, e.g., 50 for 50%): " KSM_THRES_COEF
    read -p "Enter the KSM sleep time in milliseconds (e.g., 80): " KSM_SLEEP_MSEC

    # Validate input (ensure valid numbers)
    if [[ ! "$KSM_THRES_COEF" =~ ^[0-9]+$ ]] || [[ ! "$KSM_SLEEP_MSEC" =~ ^[0-9]+$ ]]; then
        echo "Invalid input. Please enter valid numerical values."
        exit 1
    fi

    # Apply the user's input to the configuration
    sed -i -e "s/\# KSM_THRES_COEF=.*/KSM_THRES_COEF=${KSM_THRES_COEF}/g" /etc/ksmtuned.conf
    sed -i -e "s/\# KSM_SLEEP_MSEC=.*/KSM_SLEEP_MSEC=${KSM_SLEEP_MSEC}/g" /etc/ksmtuned.conf

    # Enable and start ksmtuned service
    systemctl enable ksmtuned
    systemctl start ksmtuned
    echo "KSM settings applied. KSM threshold coefficient is set to ${KSM_THRES_COEF}% and KSM sleep time is set to ${KSM_SLEEP_MSEC}ms."
fi

if [ "${PMX_AMDFIXES,,}" == "yes" ] ; then
    ## Detect AMD EPYC and Ryzen CPU and Apply Fixes
    if [ "$(grep -i -m 1 "model name" /proc/cpuinfo | grep -i "EPYC")" != "" ]; then
      echo "AMD EPYC detected"
    elif [ "$(grep -i -m 1 "model name" /proc/cpuinfo | grep -i "Ryzen")" != "" ]; then
      echo "AMD Ryzen detected"
    else
        PMX_AMDFIXES="no"
    fi

    if [ "${PMX_AMDFIXES,,}" == "yes" ] ; then
      #Apply fix to kernel : Fixes random crashing and instability
        if ! grep "GRUB_CMDLINE_LINUX_DEFAULT" /etc/default/grub | grep -q "idle=nomwait" ; then
            echo "Setting kernel idle=nomwait"
            sed -i 's/GRUB_CMDLINE_LINUX_DEFAULT="/GRUB_CMDLINE_LINUX_DEFAULT="idle=nomwait /g' /etc/default/grub
            update-grub
        fi
        ## Add msrs ignore to fix Windows guest on EPIC/Ryzen host
        echo "options kvm ignore_msrs=Y" >> /etc/modprobe.d/kvm.conf
        echo "options kvm report_ignored_msrs=N" >> /etc/modprobe.d/kvm.conf

        echo "Installing Proxmox Kernel 6.8"
        /usr/bin/env DEBIAN_FRONTEND=noninteractive apt-get -y -o Dpkg::Options::='--force-confdef' install pve-kernel-6.8
    fi
fi

if [ "${PMX_KERNELHEADERS,,}" == "yes" ] ; then
    ## Install kernel source headers
    /usr/bin/env DEBIAN_FRONTEND=noninteractive apt-get -y -o Dpkg::Options::='--force-confdef' install pve-headers module-assistant
fi

# if [ "$PMX_KEXEC" == "yes" ] ; then
# ## Install kexec, allows for quick reboots into the latest updated kernel set as primary in the boot-loader.
# # use command 'reboot-quick'
# echo "kexec-tools kexec-tools/load_kexec boolean false" | debconf-set-selections
# /usr/bin/env DEBIAN_FRONTEND=noninteractive apt-get -y -o Dpkg::Options::='--force-confdef' install kexec-tools
#     cat <<'EOF' > /etc/systemd/system/kexec-pve.service
# [Unit]
# Description=Loading new kernel into memory
# Documentation=man:kexec(8)
# DefaultDependencies=no
# Before=reboot.target
# RequiresMountsFor=/boot
# #Before=shutdown.target umount.target final.target

# [Service]
# Type=oneshot
# RemainAfterExit=yes
# ExecStart=/sbin/kexec -d -l /boot/pve/vmlinuz --initrd=/boot/pve/initrd.img --reuse-cmdline

# [Install]
# WantedBy=default.target
# EOF
#     systemctl enable kexec-pve.service
#     echo "alias reboot-quick='systemctl kexec'" >> /root/.bash_profile
# fi

if [ "${PMX_DISABLERPC,,}" == "yes" ] ; then
    ## Disable portmapper / rpcbind (Increases Security)
    echo -e "Disable portmapper / rpcbind (Increases Security."
    systemctl disable rpcbind
    systemctl stop rpcbind
fi

if [ "${PMX_TIMEZONE}" == "yes" ] ; then
    ## Set Timezone, empty = set automatically by IP
    echo -e "Set Timezone, empty = set automatically by IP"
    this_ip="$(dig +short myip.opendns.com @resolver1.opendns.com)"
    timezone="$(curl "https://ipapi.co/${this_ip}/timezone")"
    if [ "$timezone" != "" ] ; then
        echo "Found $timezone for ${this_ip}"
        timedatectl set-timezone "$timezone"
    else
        echo "WARNING: Timezone not found for ${this_ip}, set to UTC"
        timedatectl set-timezone UTC
    fi
else
    ## Set Timezone to PMX_TIMEZONE
    timedatectl set-timezone "$PMX_TIMEZONE"
fi

if [ "${PMX_TIMESYNC,,}" == "yes" ] ; then
    echo -e "Setup ntp time sync."
    timedatectl set-ntp true
fi

if [ "${PMX_PIGZ,,}" == "yes" ] ; then
    ## Set pigz to replace gzip, 2x faster gzip compression
    echo  -e "Set pigz to replace gzip, 2x faster gzip compression."
    sed -i "s/#pigz:.*/pigz: 1/" /etc/vzdump.conf
    /usr/bin/env DEBIAN_FRONTEND=noninteractive apt-get -y -o Dpkg::Options::='--force-confdef' install pigz
    cat  <<EOF > /bin/pigzwrapper
#!/bin/sh
# System Optimizations by h@x - For Proxmox
PATH=/bin:\$PATH
GZIP="-1"
exec /usr/bin/pigz "\$@"
EOF
    mv -f /bin/gzip /bin/gzip.original
    cp -f /bin/pigzwrapper /bin/gzip
    chmod +x /bin/pigzwrapper
    chmod +x /bin/gzip
fi

if [ "${PMX_FAIL2BAN,,}" == "yes" ] ; then
    ## Hardening Proxmox Web Interface with fail2ban
    echo -e "Hardening Proxmox Web Interface with fail2ban."
    /usr/bin/env DEBIAN_FRONTEND=noninteractive apt-get -y -o Dpkg::Options::='--force-confdef' install fail2ban
cat <<EOF > /etc/fail2ban/filter.d/proxmox.conf
[Definition]
failregex = pvedaemon\[.*authentication failure; rhost=<HOST> user=.* msg=.*
ignoreregex =
EOF

cat <<EOF > /etc/fail2ban/jail.d/proxmox.conf
[proxmox]
enabled = true
port = https,http,8006,8007
filter = proxmox
logpath = /var/log/daemon.log
maxretry = 3
# 1 hour
bantime = 3600
findtime = 600
EOF

# cat <<EOF > /etc/fail2ban/jail.local
# [DEFAULT]
# banaction = iptables-ipset-proto4
# EOF
systemctl enable fail2ban
#     ##testing
#     #fail2ban-regex /var/log/daemon.log /etc/fail2ban/filter.d/proxmox.conf
fi

if [ "${PMX_NOSUBBANNER,,}" == "yes" ] ; then
    ## Remove subscription banner
    echo -e "Remove Proxmox Subscription banner."
    if [ -f "/usr/share/javascript/proxmox-widget-toolkit/proxmoxlib.js" ] ; then
      # create a daily cron to make sure the banner does not re-appear
  cat <<'EOF' > /etc/cron.daily/pmx-pve-nosub
#!/bin/sh
# System Optimizations by h@x - For Proxmox Remove subscription banner
sed -i "s/data.status !== 'Active'/false/g" /usr/share/javascript/proxmox-widget-toolkit/proxmoxlib.js
sed -i "s/checked_command: function(orig_cmd) {/checked_command: function() {} || function(orig_cmd) {/g" /usr/share/javascript/proxmox-widget-toolkit/proxmoxlib.js
EOF

    fi
fi

chmod 755 /etc/cron.daily/pmx-pve-nosub
bash /etc/cron.daily/pmx-pve-nosub
# Remove nag @tinof
echo "DPkg::Post-Invoke { \"dpkg -V proxmox-widget-toolkit | grep -q '/proxmoxlib\.js$'; if [ \$? -eq 1 ]; then { echo 'Removing subscription nag from UI...'; sed -i '/data.status/{s/\!//;s/Active/NoMoreNagging/}' /usr/share/javascript/proxmox-widget-toolkit/proxmoxlib.js; }; fi\"; };" > /etc/apt/apt.conf.d/pmx-pve-no-nag && apt --reinstall install proxmox-widget-toolkit

if [ "${PMX_MOTD,,}" == "yes" ] ; then
## Pretty MOTD BANNER
echo -e "Adding pretty MOTD Banner."
  if ! grep -q https "/etc/motd" ; then
    cat << 'EOF' > /etc/motd.new
	   This system is optimised by: h@x - for Proxmox
EOF

    cat /etc/motd >> /etc/motd.new
    mv /etc/motd.new /etc/motd
  fi
fi

if [ "${PMX_KERNELPANIC,,}" == "yes" ] ; then
    # Enable restart on kernel panic
    echo -e "Enable restart on kernel panic"
    cat <<EOF > /etc/sysctl.d/99-pmx-kernelpanic.conf
# System Optimizations by h@x - For Proxmox
# Enable restart on kernel panic, kernel oops and hardlockup
kernel.core_pattern=/var/crash/core.%t.%p
# Reboot on kernel panic afetr 10s
kernel.panic=10
# Panic on kernel oops, kernel exploits generally create an oops
kernel.panic_on_oops=1
# Panic on a hardlockup
kernel.hardlockup_panic=1
EOF
fi

if [ "${PMX_LIMITS,,}" == "yes" ] ; then
    ## Increase max user watches
    # BUG FIX : No space left on device
    echo -e "Increase max user watches & Bug fix : No space left on device."
    cat <<EOF > /etc/sysctl.d/99-pmx-maxwatches.conf
# System Optimizations by h@x - For Proxmox
# Increase max user watches
fs.inotify.max_user_watches=1048576
fs.inotify.max_user_instances=1048576
fs.inotify.max_queued_events=1048576
EOF
    ## Increase max FD limit / ulimit
    echo -e "Increase max FD limit / ulimit."
    cat <<EOF >> /etc/security/limits.d/99-pmx-limits.conf
# System Optimizations by h@x - For Proxmox
# Increase max FD limit / ulimit
* soft     nproc          1048576
* hard     nproc          1048576
* soft     nofile         1048576
* hard     nofile         1048576
root soft     nproc          unlimited
root hard     nproc          unlimited
root soft     nofile         unlimited
root hard     nofile         unlimited
EOF
    ## Increase kernel max Key limit
    echo -e "Increase kernel max Key limit."
    cat <<EOF > /etc/sysctl.d/99-pmx-maxkeys.conf
# System Optimizations by h@x - For Proxmox
# Increase kernel max Key limit
kernel.keys.root_maxkeys=1000000
kernel.keys.maxkeys=1000000
EOF
    ## Set systemd ulimits
    echo -e "Increase kernel max Key limit."
    echo "DefaultLimitNOFILE=256000" >> /etc/systemd/system.conf
    echo "DefaultLimitNOFILE=256000" >> /etc/systemd/user.conf

    echo 'session required pam_limits.so' >> /etc/pam.d/common-session
    echo 'session required pam_limits.so' >> /etc/pam.d/runuser-l

    ## Set ulimit for the shell user
    echo -e "Set ulimit for the shell user."
    echo "ulimit -n 256000" >> /root/.profile
fi

if [ "${PMX_LOGROTATE,,}" == "yes" ] ; then
    ## Optimise logrotate
    echo -e "Optimise logrotate."
    cat <<EOF > /etc/logrotate.conf
# System Optimizations by h@x - For Proxmox
daily
su root adm
rotate 7
create
compress
size=10M
delaycompress
copytruncate

include /etc/logrotate.d
EOF
    systemctl restart logrotate
fi

if [ "${PMX_JOURNALD,,}" == "yes" ] ; then
    ## Limit the size and optimise journald
    echo -e "Optimise logrotate."
    cat <<EOF > /etc/systemd/journald.conf
# System Optimizations by h@x - For Proxmox
[Journal]
# Store on disk
Storage=persistent
# Don't split Journald logs by user
SplitMode=none
# Disable rate limits
RateLimitInterval=0
RateLimitIntervalSec=0
RateLimitBurst=0
# Disable Journald forwarding to syslog
ForwardToSyslog=no
# Journald forwarding to wall /var/log/kern.log
ForwardToWall=yes
# Disable signing of the logs, save cpu resources.
Seal=no
Compress=yes
# Fix the log size
SystemMaxUse=64M
RuntimeMaxUse=60M
# Optimise the logging and speed up tasks
MaxLevelStore=warning
MaxLevelSyslog=warning
MaxLevelKMsg=warning
MaxLevelConsole=notice
MaxLevelWall=crit
EOF
    systemctl restart systemd-journald.service
    journalctl --vacuum-size=64M --vacuum-time=1d;
    journalctl --rotate
fi

if [ "${PMX_ENTROPY,,}" == "yes" ] ; then
## Ensure Entropy Pools are Populated, prevents slowdowns whilst waiting for entropy
    echo -e "Ensure Entropy Pools are Populated, prevents slowdowns whilst waiting for entropy."
    /usr/bin/env DEBIAN_FRONTEND=noninteractive apt-get -y -o Dpkg::Options::='--force-confdef' install haveged
    ## Net optimising
    cat <<EOF > /etc/default/haveged
# System Optimizations by h@x - For Proxmox
#   -w sets low entropy watermark (in bits)
DAEMON_ARGS="-w 1024"
EOF
    systemctl daemon-reload
    systemctl enable haveged
fi

if [ "${PMX_VZDUMP,,}" == "yes" ] ; then
    ## Increase vzdump backup speed
    echo -e "Increasing vzdump backup speed."
    sed -i "s/#bwlimit:.*/bwlimit: 0/" /etc/vzdump.conf
    sed -i "s/#ionice:.*/ionice: 5/" /etc/vzdump.conf
fi

if [ "${PMX_MEMORYFIXES,,}" == "yes" ] ; then
    ## Optimise Memory
    echo -e "Optimising Memory."
cat <<EOF > /etc/sysctl.d/99-pmx-memory.conf
# System Optimizations by h@x - For Proxmox
# Memory Optimising
## Bugfix: reserve 1024MB memory for system
vm.min_free_kbytes=1048576
vm.nr_hugepages=72
# (Redis/MongoDB)
vm.max_map_count=262144
vm.overcommit_memory = 1
EOF
fi

if [ "${PMX_TCPBBR,,}" == "yes" ] ; then
## Enable TCP BBR congestion control
echo -e "Enable TCP BBR congestion control."
cat <<EOF > /etc/sysctl.d/99-pmx-kernel-bbr.conf
# System Optimizations by h@x - For Proxmox
# TCP BBR congestion control
net.core.default_qdisc=fq
net.ipv4.tcp_congestion_control=bbr
EOF
fi

if [ "${PMX_TCPFASTOPEN,,}" == "yes" ] ; then
## Enable TCP fastopen
echo -e "Enable TCP fastopen."
cat <<EOF > /etc/sysctl.d/99-pmx-tcp-fastopen.conf
# System Optimizations by h@x - For Proxmox
# TCP fastopen
net.ipv4.tcp_fastopen=3
EOF
fi

if [ "${PMX_NET,,}" == "yes" ] ; then
## Enable Network optimizing
echo -e "Enable Network optimizing."
cat <<EOF > /etc/sysctl.d/99-pmx-net.conf
# System Optimizations by h@x - For Proxmox
net.core.netdev_max_backlog=8192
net.core.optmem_max=8192
net.core.rmem_max=16777216
net.core.somaxconn=8151
net.core.wmem_max=16777216
net.ipv4.conf.all.accept_redirects = 0
net.ipv4.conf.all.accept_source_route = 0
net.ipv4.conf.all.log_martians = 0
net.ipv4.conf.all.rp_filter = 1
net.ipv4.conf.all.secure_redirects = 0
net.ipv4.conf.all.send_redirects = 0
net.ipv4.conf.default.accept_redirects = 0
net.ipv4.conf.default.accept_source_route = 0
net.ipv4.conf.default.log_martians = 0
net.ipv4.conf.default.rp_filter = 1
net.ipv4.conf.default.secure_redirects = 0
net.ipv4.conf.default.send_redirects = 0
net.ipv4.icmp_echo_ignore_broadcasts = 1
net.ipv4.icmp_ignore_bogus_error_responses = 1
net.ipv4.ip_local_port_range=1024 65535
net.ipv4.tcp_base_mss = 1024
net.ipv4.tcp_challenge_ack_limit = 999999999
net.ipv4.tcp_fin_timeout=10
net.ipv4.tcp_keepalive_intvl=30
net.ipv4.tcp_keepalive_probes=3
net.ipv4.tcp_keepalive_time=240
net.ipv4.tcp_limit_output_bytes=65536
net.ipv4.tcp_max_syn_backlog=8192
net.ipv4.tcp_max_tw_buckets = 1440000
net.ipv4.tcp_mtu_probing = 1
net.ipv4.tcp_rfc1337=1
net.ipv4.tcp_rmem=8192 87380 16777216
net.ipv4.tcp_sack=1
net.ipv4.tcp_slow_start_after_idle=0
net.ipv4.tcp_syn_retries=3
net.ipv4.tcp_synack_retries = 2
net.ipv4.tcp_tw_recycle = 0
net.ipv4.tcp_tw_reuse = 0
net.ipv4.tcp_wmem=8192 65536 16777216
net.netfilter.nf_conntrack_generic_timeout = 60
net.netfilter.nf_conntrack_helper=0
net.netfilter.nf_conntrack_max = 524288
net.netfilter.nf_conntrack_tcp_timeout_established = 28800
net.unix.max_dgram_qlen = 4096
EOF
fi

if [ "${PMX_SWAPPINESS,,}" == "yes" ] ; then
    ## Bugfix: high swap usage with low memory usage
    echo -e "Bugfix: high swap usage with low memory usage."
    cat <<EOF > /etc/sysctl.d/99-pmx-swap.conf
# System Optimizations by h@x - For Proxmox
# Bugfix: high swap usage with low memory usage
vm.swappiness=10
EOF
fi

if [ "${PMX_MAXFS,,}" == "yes" ] ; then
    ## Increase Max FS open files
    echo -e "Increase Max FS open files."
    cat <<EOF > /etc/sysctl.d/99-pmx-fs.conf
# System Optimizations by h@x - For Proxmox
# Max FS Optimising
fs.nr_open=12000000
fs.file-max=9000000
fs.aio-max-nr=524288
EOF
fi

if [ "${PMX_BASHRC,,}" == "yes" ] ; then
    ## Customize bashrc
    cat <<EOF >> /root/.bashrc
export HISTTIMEFORMAT="%d/%m/%y %T "
export PS1='\u@\h:\W \$ '
alias l='ls -CF'
alias la='ls -A'
alias ll='ls -alF'
alias ls='ls --color=auto'
source /etc/profile.d/bash_completion.sh
export PS1="\[\e[31m\][\[\e[m\]\[\e[38;5;172m\]\u\[\e[m\]@\[\e[38;5;153m\]\h\[\e[m\] \[\e[38;5;214m\]\W\[\e[m\]\[\e[31m\]]\[\e[m\]\\$ "
EOF
    echo "source /root/.bashrc" >> /root/.bash_profile
fi

if [ "${PMX_ZFSARC,,}" == "yes" ] ; then
    ## Optimise ZFS arc size accoring to memory size
    echo -e "Optimise ZFS arc size accoring to memory size."
    if [ "$(command -v zfs)" != "" ] ; then
      if [[ RAM_SIZE_GB -le 16 ]] ; then
        MY_ZFS_ARC_MIN=536870911
        MY_ZFS_ARC_MAX=536870912
    elif [[ RAM_SIZE_GB -le 32 ]] ; then
        # 1GB/1GB
        MY_ZFS_ARC_MIN=1073741823
        MY_ZFS_ARC_MAX=1073741824
      else
        MY_ZFS_ARC_MIN=$((RAM_SIZE_GB * 1073741824 / 16))
        MY_ZFS_ARC_MAX=$((RAM_SIZE_GB * 1073741824 / 8))
      fi
      # Enforce the minimum, incase of a faulty vmstat
      if [[ MY_ZFS_ARC_MIN -lt 536870911 ]] ; then
        echo -e "Enforce the minimum, incase of a faulty vmstat".
        MY_ZFS_ARC_MIN=536870911
      fi
      if [[ MY_ZFS_ARC_MAX -lt 536870912 ]] ; then
        MY_ZFS_ARC_MAX=536870912
      fi
      cat <<EOF > /etc/modprobe.d/99-pmx-zfsarc.conf
# System Optimizations by h@x - For Proxmox ZFS tuning

# Use 1/8 RAM for MAX cache, 1/16 RAM for MIN cache, or 1GB
options zfs zfs_arc_min=$MY_ZFS_ARC_MIN
options zfs zfs_arc_max=$MY_ZFS_ARC_MAX

# use the prefetch method
options zfs l2arc_noprefetch=0

# max write speed to l2arc
# tradeoff between write/read and durability of ssd (?)
# default : 8 * 1024 * 1024
# setting here : 500 * 1024 * 1024
options zfs l2arc_write_max=524288000
options zfs zfs_txg_timeout=60
EOF
    fi
fi


# Fix missing /etc/network/interfaces.d include
echo -e "Fix missing /etc/network/interfaces.d include."
if ! grep -q 'source /etc/network/interfaces.d/*' "/etc/network/interfaces" ; then
    echo "Added missing include to /etc/network/interfaces"
    echo "source /etc/network/interfaces.d/*" >> /etc/network/interfaces
fi

if [ "${PMX_VFIO_IOMMU,,}" == "yes" ] ; then
    # Enable IOMMU
    echo -e "Enable IOMMU."
    cpu=$(cat /proc/cpuinfo)
    if [[ $cpu == *"GenuineIntel"* ]]; then
        echo "Detected Intel CPU"
        sed -i 's/quiet/quiet intel_iommu=on iommu=pt/g' /etc/default/grub
    elif [[ $cpu == *"AuthenticAMD"* ]]; then
        echo "Detected AMD CPU"
        sed -i 's/quiet/quiet amd_iommu=on iommu=pt/g' /etc/default/grub
    else
        echo "Unknown CPU! Cannot enable IOMMU!"
    fi

    cat <<EOF >> /etc/modules
# System Optimizations by h@x - For Proxmox
vfio
vfio_iommu_type1
vfio_pci
vfio_virqfd

EOF
echo -e "Add common gpu drivers on a blacklist."
cat <<EOF >> /etc/modprobe.d/blacklist.conf
# System Optimizations by h@x - For Proxmox
blacklist nouveau
blacklist lbm-nouveau
options nouveau modeset=0
blacklist amdgpu
blacklist radeon
blacklist nvidia
blacklist nvidiafb

EOF
fi

# propagate the settings
echo -e "Apply changes."
update-initramfs -u -k all
update-grub
pve-efiboot-tool refresh

# cleanup
## Remove no longer required packages and purge old cached updates
echo -e "Remove no longer required packages and purge old cached updates"
/usr/bin/env DEBIAN_FRONTEND=noninteractive apt-get -y -o Dpkg::Options::='--force-confdef' autoremove
/usr/bin/env DEBIAN_FRONTEND=noninteractive apt-get -y -o Dpkg::Options::='--force-confdef' autoclean

echo "# System optimization has been installed." > /etc/system-optimization
date >> /etc/system-optimization
}

function finish() {
## Script Finish
echo -e '\033[1;33m Finished....please restart the system \033[0m'
echo "Optimizations by h@x"
}

function check_root() {
    if [ "$(id -u)" -ne 0 ]; then
        echo "Error: This script must be run as root."
        exit 1
    fi
}

check-root
ask-user
#main
#finish