install_nvidia_drivers.sh

#!/bin/bash

# ==============================================================================
# NVIDIA Driver Installation Script
#
# This script automates the installation of NVIDIA drivers on various Linux
# distributions based on the official NVIDIA documentation:
# https://docs.nvidia.com/datacenter/tesla/driver-installation-guide/index.html
#
# Supported Distributions:
# - Debian-based: Debian 12/13, Ubuntu 22.04, Ubuntu 24.04
# - RHEL-based: RHEL 8/9, Rocky Linux 8/9, Oracle Linux 8/9, Fedora
# - SUSE-based: SUSE Linux Enterprise 15, openSUSE Leap 15
#
# ==============================================================================

set -e

# --- Helper Functions ---

_print_info() {
    echo -e "\e[34m[INFO]\e[0m $1"
}

_print_error() {
    echo -e "\e[31m[ERROR]\e[0m $1" >&2
    exit 1
}

# --- Installation Functions ---

install_dependencies_apt() {
    _print_info "Installing kernel headers and dependencies for Debian/Ubuntu..."
    apt-get update
    apt-get install -y linux-headers-$(uname -r)
}

install_dependencies_dnf() {
    _print_info "Installing kernel headers and dependencies for RHEL/Fedora..."
    if [[ "$ID" == "rhel" || "$ID" == "rocky" || "$ID" == "ol" ]]; then
        if [[ "$VERSION_ID" == "8" ]]; then
            dnf install -y "kernel-devel-$(uname -r)" kernel-headers
        else
            dnf install -y kernel-devel-matched kernel-headers
        fi
        dnf config-manager --set-enabled crb || dnf config-manager --set-enabled powertools || dnf config-manager --set-enabled ol${VERSION_ID}_codeready_builder
        dnf install -y epel-release || dnf install -y oracle-epel-release-el${VERSION_ID}
    elif [[ "$ID" == "fedora" ]]; then
        dnf install -y kernel-devel-matched kernel-headers
    fi
}

install_dependencies_zypper() {
    _print_info "Installing kernel headers and dependencies for SUSE..."
    zypper install -y kernel-default-devel=$(uname -r | sed 's/\-default//')
}

add_repo_apt() {
    local codename=$1
    local arch=$(dpkg --print-architecture)
    if [ "$arch" = "amd64" ]; then
        arch="x86_64"
    fi

    _print_info "Adding NVIDIA repository for Debian/Ubuntu..."
    wget https://developer.download.nvidia.com/compute/cuda/repos/$codename/$arch/cuda-keyring_1.1-1_all.deb
    dpkg -i cuda-keyring_1.1-1_all.deb
    apt-get update
    rm cuda-keyring_1.1-1_all.deb
}

add_repo_dnf() {
    local codename=$1
    local arch=$(uname -m)
    _print_info "Adding NVIDIA repository for RHEL/Fedora..."
    if [ "$arch" = "aarch64" ]; then
        dnf config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/$codename/sbsa/cuda-$codename.repo
    else
        dnf config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/$codename/$arch/cuda-$codename.repo
    fi
    dnf clean expire-cache
}

add_repo_zypper() {
    local codename=$1
    local arch=$(uname -m)
    _print_info "Adding NVIDIA repository for SUSE..."
    if [ "$arch" = "aarch64" ]; then
        zypper addrepo https://developer.download.nvidia.com/compute/cuda/repos/$codename/sbsa/cuda-$codename.repo
    else
        zypper addrepo https://developer.download.nvidia.com/compute/cuda/repos/$codename/$arch/cuda-$codename.repo
    fi
    SUSEConnect --product PackageHub/15/$(uname -m) || true
    zypper refresh
}

install_driver_apt() {
    _print_info "Installing NVIDIA drivers (cuda-drivers) for Debian/Ubuntu..."
    apt-get install -y cuda-drivers
}

install_driver_dnf() {
    _print_info "Installing NVIDIA drivers (cuda-drivers) for RHEL/Fedora..."
    dnf install -y cuda-drivers
}

install_driver_zypper() {
    _print_info "Installing NVIDIA drivers (cuda-drivers) for SUSE..."
    zypper -n install cuda-drivers
}

# --- Platform-Specific Main Functions ---

install_debian_based() {
    local codename
    case "$ID-$VERSION_ID" in
        "ubuntu-22.04") codename="ubuntu2204" ;;
        "ubuntu-24.04") codename="ubuntu2404" ;;
        "debian-12") codename="debian12" ;;
        "debian-13") codename="debian13" ;;
        *) _print_error "Unsupported Debian/Ubuntu version: $ID $VERSION_ID" ;;
    esac

    install_dependencies_apt
    add_repo_apt "$codename"
    install_driver_apt
}

install_rhel_based() {
    local codename
    case "$ID-$VERSION_ID" in
        "rhel-8" | "rocky-8" | "ol-8") codename="rhel8" ;;
        "rhel-9" | "rocky-9" | "ol-9") codename="rhel9" ;;
        "fedora-42") codename="fedora42" ;;
        *) _print_error "Unsupported RHEL/Fedora version: $ID $VERSION_ID" ;;
    esac
    
    install_dependencies_dnf
    add_repo_dnf "$codename"
    install_driver_dnf
}

install_suse_based() {
    local codename
    case "$ID-$VERSION_ID" in
        "sles-15"*) codename="sles15" ;;
        "opensuse-leap-15"*) codename="opensuse15" ;;
        *) _print_error "Unsupported SUSE version: $ID $VERSION_ID" ;;
    esac
    
    install_dependencies_zypper
    add_repo_zypper "$codename"
    install_driver_zypper
}

# --- Main Script ---

main() {
    _print_info "Starting NVIDIA Driver Installation Script"

    if [ "$EUID" -ne 0 ]; then
        _print_error "This script must be run as root."
    fi

    if [ ! -f /etc/os-release ]; then
        _print_error "Cannot detect OS distribution."
    fi

    # Source os-release to get distribution info
    . /etc/os-release

    _print_info "Detected Distribution: $PRETTY_NAME"

    case "$ID_LIKE" in
        *debian*)
            install_debian_based
            ;;
        *rhel*|*fedora*)
            install_rhel_based
            ;;
        *suse*)
            install_suse_based
            ;;
        *)
            # Fallback for systems that don't set ID_LIKE
            case "$ID" in
                ubuntu|debian)
                    install_debian_based
                    ;;
                rhel|rocky|ol|fedora)
                    install_rhel_based
                    ;;
                sles|opensuse-leap)
                    install_suse_based
                    ;;
                *)
                    _print_error "Unsupported distribution: $ID"
                    ;;
            esac
            ;;
    esac

    _print_info "NVIDIA driver installation completed successfully."
    _print_info "A reboot is required for the changes to take effect."
    read -p "Reboot now? (y/n): " REBOOT_CHOICE
    if [[ "$REBOOT_CHOICE" == "y" || "$REBOOT_CHOICE" == "Y" ]]; then
        reboot
    fi
}

main "$@"