#!/bin/bash
# rbt

BOLD=$(tput bold)
NORMAL=$(tput sgr0)
GREY="\e[2m"
reboot_check=

# На будущее
if [ "${USER}" != "root" ]; then
    echo -e "Выполните используя sudo:\nsudo $(basename "$0")"
    exit -1
fi

function install_depends() {
    local depends=($*)
    local not_installed=()
    for DN in ${depends[*]}; do
        dpkg -s ${DN} &> /dev/null
        [ $? -ne 0 ] && not_installed+=(${DN})
    done
    if [ -n "$not_installed" ]; then
        apt update &> /dev/null
        if [ $? -ne 0 ]; then
            echo "E: Не удалось обновить список доступных пакетов"
            exit -1
        fi
        echo -e "\n${BOLD}Установка необходимых пакетов${NORMAL}\n"
        apt install -y ${not_installed[*]}
        [ $? -ne 0 ] && exit -1
        echo
    fi
}

function edit_grub() {
    local grub_path=/etc/default/grub
    local def_string= def_par= def_other_par= new_string=
    def_string=$(awk '!/^#/ && /GRUB_CMDLINE_LINUX_DEFAULT="/ {print}' $grub_path)
    def_par=$(awk '{gsub(/(.*GRUB_CMDLINE_LINUX_DEFAULT=)|(\")/,"");print}' <<< "$def_string")
    if ! grep -wq 'intel_iommu *= *on' <<< ${def_string}; then
        def_other_par=$(awk '{gsub(/( *intel_iommu *= *[a-z0-1]*)/,"");print}' <<< "$def_par")
        # new_string="$(echo GRUB_CMDLINE_LINUX_DEFAULT=\"${def_other_par} intel_iommu=on\")"
        new_string=$(tr -s ' ' <<< "GRUB_CMDLINE_LINUX_DEFAULT=\"${def_other_par} intel_iommu=on\"")
        sed -i "/${def_string}/c  # ${def_string}\n${new_string}" $grub_path
        echo -e "${GREY}\n${grub_path}\nСтрока:\t\t${def_string}\nИзменена на:\t${new_string}${NORMAL}\n"
        update-grub
        [ $? -eq 0 ] && reboot_check=X
    else
        if ! $(dmesg | grep -Ewq "DMAR: *IOMMU"); then
            reboot_check=X
        fi
    fi
}

function edit_blacklist_conf() {
   local blacklist_conf=/etc/modprobe.d/blacklist.conf
   local blacklist_nouveau='blacklist nouveau'
    if ! grep -Ewq "^ *blacklist +nouveau" $blacklist_conf; then
        echo -e "\n#rbt: for vgpu\n${blacklist_nouveau}" >> $blacklist_conf
        echo -e "\n${GREY}${blacklist_conf}\nДобавлена строка: ${blacklist_nouveau}${NORMAL}\n"
        reboot_check=X
    fi
}

function edit_opennebula_vfio_rules() {
    local opennebula_vfio_rules=/etc/udev/rules.d/opennebula-vfio.rules
    local udev_rules='SUBSYSTEM=="vfio", OWNER="oneadmin", GROUP="oneadmin", MODE="0600"'
    if [ ! -f ${opennebula_vfio_rules} ] || ! grep -Ewq 'SUBSYSTEM=="vfio", *OWNER="oneadmin", *GROUP="oneadmin", *MODE="0600" *$' ${opennebula_vfio_rules}; then
        echo "${udev_rules}" > ${opennebula_vfio_rules}
        echo -e "\n${GREY}${opennebula_vfio_rules}\nДобавлена строка: ${udev_rules}${NORMAL}\n"
        udevadm control --reload-rules && udevadm trigger
        reboot_check=X
    fi
}

function nvidia_sriov_service(){
    local nvidia_sriov_service=/etc/systemd/system/nvidia-sriov.service
    local NVIDIA_SRIOV_SERVICE="$(cat << EOF
[Unit]
Description=Enable NVIDIA SR-IOV
After=network.target nvidia-vgpud.service nvidia-vgpu-mgr.service
Before=pve-guests.service
[Service]
Type=oneshot
ExecStart=/usr/lib/nvidia/sriov-manage -e ALL
[Install]
WantedBy=multi-user.target
EOF
)"
    if ! [ -f ${nvidia_sriov_service} ]; then
        echo "${NVIDIA_SRIOV_SERVICE}" > ${nvidia_sriov_service}
        systemctl daemon-reload
    fi
    systemctl enable nvidia-sriov.service
}

# Install packages:
install_depends gcc make

# Enable Memory Management Unit for I/O Operations IOMMU
edit_grub

# Disable nouveau module
edit_blacklist_conf

# Create and apply a udev rule
edit_opennebula_vfio_rules

# Create service nvidia-sriov.service
nvidia_sriov_service

if [ -n "$reboot_check" ]; then
    echo -e "\n\t${BOLD}Необходимо перезагрузить систему!${NORMAL}\n"
fi
