#!/bin/bash

# ==============================================================================
# Script Name: hpc-security-hardener.sh
# Description: Mitigates CopyFail (AF_ALG) and DirtyFrag (ESP/RXRPC) vulnerabilities.
# Architect Note: This script ensures idempotency for HPC cluster deployment.
# ==============================================================================

set -euo pipefail

# --- Professional Logging ---
log_info() { echo -e "\033[32m[INFO]\033[0m $1"; }
log_warn() { echo -e "\033[33m[WARN]\033[0m $1"; }

# --- Ensure Root Privileges ---
if [[ $EUID -ne 0 ]]; then
   echo "[ERROR] This script must be run as root"
   exit 1
fi

# --- Function: Fix CopyFail (AF_ALG) ---
fix_copy_fail() {
    log_info "Starting CopyFail (AF_ALG) mitigation..."

    local mod_conf="/etc/modprobe.d/disable-algif-aead.conf"
    local blacklist_args="initcall_blacklist=algif_aead_init,af_alg_init"

    # 1. Update Kernel Args via grubby (Idempotent check)
    # Check if the arguments already exist in the current default kernel info
    if ! grubby --info=DEFAULT | grep -q "initcall_blacklist=algif_aead_init"; then
        log_info "Appending initcall_blacklist to kernel arguments..."
        sudo grubby --update-kernel=ALL --args="$blacklist_args"
    else
        log_info "Kernel arguments for CopyFail already present. Skipping."
    fi

    # 2. Disable module loading via modprobe
    if [[ ! -f "$mod_conf" ]] || ! grep -q "algif_aead" "$mod_conf"; then
        log_info "Creating $mod_conf..."
        echo "install algif_aead /bin/false" > "$mod_conf"
    else
        log_info "Module algif_aead already disabled in modprobe. Skipping."
    fi

    # 3. Attempt to unload module if currently loaded
    if lsmod | grep -q "algif_aead"; then
        log_info "Unloading algif_aead module..."
        rmmod algif_aead 2>/dev/null || log_warn "Failed to unload algif_aead (it might be in use)."
    fi
}

# --- Function: Fix DirtyFrag (ESP/RXRPC) ---
fix_dirty_frag() {
    log_info "Starting DirtyFrag (ESP/RXRPC) mitigation..."

    local mod_conf="/etc/modprobe.d/dirtyfrag.conf"
    local config_content="install esp4 /bin/false\ninstall esp6 /bin/false\ninstall rxrpc /bin/false"

    # 1. Disable modules via modprobe
    if [[ ! -f "$mod_conf" ]]; then
        log_info "Creating $mod_conf..."
        printf "$config_content\n" > "$mod_conf"
    else
        log_info "DirtyFrag modprobe config already exists. Skipping."
    fi

    # 2. Attempt to unload modules
    for mod in esp4 esp6 rxrpc; do
        if lsmod | grep -q "^$mod"; then
            log_info "Unloading $mod module..."
            rmmod "$mod" 2>/dev/null || log_warn "Failed to unload $mod (it might be in use)."
        fi
    done
}

# --- Function: Memory Sanitization ---
clear_caches() {
    log_info "Dropping system caches to ensure memory purity..."
    sync
    echo 3 > /proc/sys/vm/drop_caches
}

# --- Main Execution ---
main() {
    log_info "HPC Security Hardening initiated on $(hostname)"

    fix_copy_fail
    fix_dirty_frag
    clear_caches

    log_info "Mitigation complete. Note: Some changes (grubby) require a reboot to take effect."
}

main "$@"