From 5db4456dba9f8c85ff70e545fa5bb3d6a37bd265 Mon Sep 17 00:00:00 2001
From: codefiles <11915375+codefiles@users.noreply.github.com>
Date: Fri, 1 Dec 2023 09:59:44 -0500
Subject: Restrict UKI to supported options (#2280)

---
 archinstall/__init__.py              |  3 +++
 archinstall/lib/global_menu.py       | 14 +++++++++-----
 archinstall/lib/installer.py         | 16 +++++++---------
 archinstall/lib/models/bootloader.py |  7 +++++++
 4 files changed, 26 insertions(+), 14 deletions(-)

diff --git a/archinstall/__init__.py b/archinstall/__init__.py
index 59b8e6b2..2625df67 100644
--- a/archinstall/__init__.py
+++ b/archinstall/__init__.py
@@ -244,6 +244,9 @@ def load_config():
 	if arguments.get('bootloader', None) is not None:
 		arguments['bootloader'] = models.Bootloader.from_arg(arguments['bootloader'])
 
+		if arguments.get('uki') and not arguments['bootloader'].has_uki_support():
+			arguments['uki'] = False
+
 	if arguments.get('audio_config', None) is not None:
 		arguments['audio_config'] = models.AudioConfiguration.parse_arg(arguments['audio_config'])
 
diff --git a/archinstall/lib/global_menu.py b/archinstall/lib/global_menu.py
index d3d87603..e65915db 100644
--- a/archinstall/lib/global_menu.py
+++ b/archinstall/lib/global_menu.py
@@ -218,14 +218,20 @@ class GlobalMenu(AbstractMenu):
 			return False
 		return self._validate_bootloader() is None
 
+	def _update_uki_display(self, name: Optional[str] = None):
+		if bootloader := self._menu_options['bootloader'].current_selection:
+			if not SysInfo.has_uefi() or not bootloader.has_uki_support():
+				self._menu_options['uki'].set_current_selection(False)
+				self._menu_options['uki'].set_enabled(False)
+			elif name and name == 'bootloader':
+				self._menu_options['uki'].set_enabled(True)
+
 	def _update_install_text(self, name: Optional[str] = None, value: Any = None):
 		text = self._install_text()
 		self._menu_options['install'].update_description(text)
 
 	def post_callback(self, name: Optional[str] = None, value: Any = None):
-		if not SysInfo.has_uefi():
-			self._menu_options['uki'].set_enabled(False)
-
+		self._update_uki_display(name)
 		self._update_install_text(name, value)
 
 	def _install_text(self):
@@ -366,8 +372,6 @@ class GlobalMenu(AbstractMenu):
 		if bootloader == Bootloader.Limine:
 			if boot_partition.fs_type != disk.FilesystemType.Fat32:
 				return "Limine does not support booting from filesystems other than FAT32"
-			elif self._menu_options['uki'].current_selection:
-				return "Limine does not support booting UKIs"
 
 		return None
 
diff --git a/archinstall/lib/installer.py b/archinstall/lib/installer.py
index d39eb922..ccae8faa 100644
--- a/archinstall/lib/installer.py
+++ b/archinstall/lib/installer.py
@@ -907,19 +907,17 @@ class Installer:
 		self,
 		boot_partition: disk.PartitionModification,
 		root_partition: disk.PartitionModification,
-		efi_partition: Optional[disk.PartitionModification],
-		uki_enabled: bool = False
+		efi_partition: Optional[disk.PartitionModification]
 	):
 		self.pacman.strap('grub')  # no need?
 
-		if not uki_enabled:
-			grub_default = self.target / 'etc/default/grub'
-			config = grub_default.read_text()
+		grub_default = self.target / 'etc/default/grub'
+		config = grub_default.read_text()
 
-			kernel_parameters = ' '.join(self._get_kernel_params(root_partition, False, False))
-			config = re.sub(r'(GRUB_CMDLINE_LINUX=")("\n)', rf'\1{kernel_parameters}\2', config, 1)
+		kernel_parameters = ' '.join(self._get_kernel_params(root_partition, False, False))
+		config = re.sub(r'(GRUB_CMDLINE_LINUX=")("\n)', rf'\1{kernel_parameters}\2', config, 1)
 
-			grub_default.write_text(config)
+		grub_default.write_text(config)
 
 		info(f"GRUB boot partition: {boot_partition.dev_path}")
 
@@ -1234,7 +1232,7 @@ Exec = /bin/sh -c "{hook_command}"
 			case Bootloader.Systemd:
 				self._add_systemd_bootloader(boot_partition, root_partition, efi_partition, uki_enabled)
 			case Bootloader.Grub:
-				self._add_grub_bootloader(boot_partition, root_partition, efi_partition, uki_enabled)
+				self._add_grub_bootloader(boot_partition, root_partition, efi_partition)
 			case Bootloader.Efistub:
 				self._add_efistub_bootloader(boot_partition, root_partition, uki_enabled)
 			case Bootloader.Limine:
diff --git a/archinstall/lib/models/bootloader.py b/archinstall/lib/models/bootloader.py
index fa3f32c8..aa1a8e27 100644
--- a/archinstall/lib/models/bootloader.py
+++ b/archinstall/lib/models/bootloader.py
@@ -14,6 +14,13 @@ class Bootloader(Enum):
 	Efistub = 'Efistub'
 	Limine = 'Limine'
 
+	def has_uki_support(self) -> bool:
+		match self:
+			case Bootloader.Efistub | Bootloader.Systemd:
+				return True
+			case _:
+				return False
+
 	def json(self) -> str:
 		return self.value
 
-- 
cgit v1.2.3-70-g09d2