#!/bin/bash
# setup_nvidia_docker.sh
# Prepares Ubuntu 24.04 host for GPU-enabled Docker containers with Hashcat

set -euo pipefail

echo "[*] Updating system..."
apt update && apt -y upgrade

echo "[*] Installing prerequisites..."
apt -y install curl wget gnupg software-properties-common lsb-release

echo "[*] Checking for NVIDIA driver..."
if ! command -v nvidia-smi &>/dev/null; then
  echo "[*] Installing NVIDIA driver..."
  apt -y install nvidia-driver-575 nvidia-utils-575
else
  echo "[+] NVIDIA driver already installed."
  nvidia-smi || true
fi

echo "[*] Installing Docker..."
if ! command -v docker &>/dev/null; then
  curl -fsSL https://get.docker.com | sh
  systemctl enable --now docker
else
  echo "[+] Docker already installed."
fi

echo "[*] Installing Docker Buildx plugin..."
apt -y install docker-buildx-plugin || {
  echo "[!] docker-buildx-plugin not available via apt, installing manually..."
  mkdir -p /usr/local/lib/docker/cli-plugins
  BUILDX_VERSION=$(curl -s https://api.github.com/repos/docker/buildx/releases/latest | grep '"tag_name"' | cut -d'"' -f4)
  curl -Lo /usr/local/lib/docker/cli-plugins/docker-buildx \
    "https://github.com/docker/buildx/releases/download/${BUILDX_VERSION}/buildx-${BUILDX_VERSION}.linux-amd64"
  chmod +x /usr/local/lib/docker/cli-plugins/docker-buildx
}

echo "[*] Installing NVIDIA Container Toolkit..."
# Remove any corrupted source list from previous runs
rm -f /etc/apt/sources.list.d/nvidia-container-toolkit.list
rm -f /etc/apt/sources.list.d/nvidia-toolkit.list
rm -f /usr/share/keyrings/nvidia-container-toolkit.gpg
rm -f /usr/share/keyrings/nvidia-toolkit.gpg

# Install GPG key
curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | \
  gpg --dearmor -o /usr/share/keyrings/nvidia-toolkit.gpg

# Download the official source list and add signed-by
curl -fsSL https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | \
  tee /etc/apt/sources.list.d/nvidia-toolkit.list
sed -i -e "s/^deb/deb [signed-by=\/usr\/share\/keyrings\/nvidia-toolkit.gpg]/g" /etc/apt/sources.list.d/nvidia-toolkit.list

apt update
apt -y install nvidia-container-toolkit

echo "[*] Configuring Docker runtime..."
nvidia-ctk runtime configure --runtime=docker
systemctl restart docker

echo
echo "[+] Setup complete!"
echo "[+] Host will require a restart before Continuing"
echo "    Test GPU passthrough with:"
echo "    docker run --rm --gpus all nvidia/cuda:12.4.1-base-ubuntu22.04 nvidia-smi"
echo