From 87c695e37f62b2f9da6f37726e31784f5bbd6243 Mon Sep 17 00:00:00 2001
From: mirkobrombin <send@mirko.pm>
Date: Thu, 1 Aug 2024 19:45:28 +0200
Subject: [PATCH] feat: support new nvidia images

---
 vanilla_installer/defaults/nvidia.py | 42 ++++++++++++++++++++++++----
 vanilla_installer/utils/processor.py |  4 ++-
 2 files changed, 39 insertions(+), 7 deletions(-)

diff --git a/vanilla_installer/defaults/nvidia.py b/vanilla_installer/defaults/nvidia.py
index c783cf95..9aa53ba4 100644
--- a/vanilla_installer/defaults/nvidia.py
+++ b/vanilla_installer/defaults/nvidia.py
@@ -16,6 +16,9 @@
 
 from gi.repository import Adw, Gtk
 
+import subprocess
+
+
 @Gtk.Template(resource_path="/org/vanillaos/Installer/gtk/default-nvidia.ui")
 class VanillaDefaultNvidia(Adw.Bin):
     __gtype_name__ = "VanillaDefaultNvidia"
@@ -27,6 +30,8 @@ class VanillaDefaultNvidia(Adw.Bin):
     info_popover = Gtk.Template.Child()
 
     use_proprietary = None
+    use_open = None
+    suggested_drivers = None
 
     def __init__(self, window, distro_info, key, step, **kwargs):
         super().__init__(**kwargs)
@@ -36,24 +41,49 @@ def __init__(self, window, distro_info, key, step, **kwargs):
         self.__step = step
         self.delta = False
 
-        self.btn_yes.connect("clicked", self.use_proprietary_drivers)
-        self.btn_no.connect("clicked", self.use_open_drivers)
+        self.suggested_drivers = self.get_suggested_drivers()
+
+        if self.suggested_drivers == "open":
+            self.info_popover.set_visible(False)
+
+        self.btn_yes.connect("clicked", self.use_drivers)
+        self.btn_no.connect("clicked", self.no_drivers)
         self.btn_info.connect("clicked", self.show_info_popover)
 
     def get_finals(self):
         return {
             "nvidia": {
                 "use-proprietary": self.use_proprietary,
+                "use-open": self.use_open,
             }
         }
 
-    def use_open_drivers(self, _):
-        self.use_proprietary = False
+    def use_drivers(self, _):
+        self.use_open = self.suggested_drivers == "open"
+        self.use_proprietary = self.suggested_drivers == "proprietary"
+
         self.__window.next()
 
-    def use_proprietary_drivers(self, _):
-        self.use_proprietary = True
+    def no_drivers(self, _):
+        self.use_proprietary = False
+        self.use_open = False
         self.__window.next()
 
     def show_info_popover(self, _):
         self.info_popover.popup()
+
+    def get_suggested_drivers(self):
+        """
+        lspci | grep -E 'NVIDIA.*(GeForce [4-8][0-9]{2}|GeForce GTX [6-9]..)' && exit 0 || exit 1
+        0 is legacy, 1 is non-legacy
+        """
+        res = subprocess.run(
+            "lspci | grep -E 'NVIDIA.*(GeForce [4-8][0-9]{2}|GeForce GTX [6-9]..)' && exit 0 || exit 1",
+            shell=True,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+        )
+        if res.returncode == 0:
+            return "proprietary"
+
+        return "open"
diff --git a/vanilla_installer/utils/processor.py b/vanilla_installer/utils/processor.py
index 9bd3b0ee..1e45c4b1 100644
--- a/vanilla_installer/utils/processor.py
+++ b/vanilla_installer/utils/processor.py
@@ -487,6 +487,8 @@ def gen_install_recipe(log_path, finals, sys_recipe):
                 oci_image = final["custom_image"]
             elif "nvidia" in final.keys():
                 if final["nvidia"]["use-proprietary"]:
+                    oci_image = images["nvidia-propietary"]
+                if final["nvidia"]["use-open"]:
                     oci_image = images["nvidia"]
             elif "vm" in final.keys():
                 if final["vm"]["use-vm-tools"]:
@@ -722,7 +724,7 @@ def gen_install_recipe(log_path, finals, sys_recipe):
             [
                 "mkdir -p /etc/abroot",
                 'echo "$(head -n-1 /usr/share/abroot/abroot.json),\n  \\"PartCryptVar\\": \\"/dev/mapper/vos--var-var\\",\n  \\"thinProvisioning\\": true,\n    \\"thinInitVolume\\": \\"vos-init\\"\n}" > /etc/abroot/abroot.json',
-                'cp /etc/abroot/abroot.json /usr/share/abroot/abroot.json',
+                "cp /etc/abroot/abroot.json /usr/share/abroot/abroot.json",
             ],
             chroot=True,
         )