Global menu (#806)

* Global menu

* Fix flake8

* Refactor code

* Add documentation

* Fix flake8

* Add support for user flow mentioned in #799

* Move import

* Fix flake8 (again)

Co-authored-by: Daniel Girtler <girtler.daniel@gmail.com>
Co-authored-by: Anton Hvornum <anton.feeds@gmail.com>

6 files changed, 496 insertions, 12 deletions

diff --git a/archinstall/lib/locale_helpers.py b/archinstall/lib/locale_helpers.py
index 6aa678a6..cbba8d52 100644
--- a/archinstall/lib/locale_helpers.py
+++ b/archinstall/lib/locale_helpers.py
@@ -1,5 +1,5 @@
 import logging
-from typing import Iterator
+from typing import Iterator, List
 
 from .exceptions import ServiceException
 from .general import SysCommand
@@ -11,6 +11,24 @@ def list_keyboard_languages() -> Iterator[str]:
 		yield line.decode('UTF-8').strip()
 
 
+def list_locales() -> List[str]:
+	with open('/etc/locale.gen', 'r') as fp:
+		locales = []
+		# before the list of locales begins there's an empty line with a '#' in front
+		# so we'll collect the localels from bottom up and halt when we're donw
+		entries = fp.readlines()
+		entries.reverse()
+
+		for entry in entries:
+			text = entry[1:].strip()
+			if text == '':
+				break
+			locales.append(text)
+
+		locales.reverse()
+		return locales
+
+
 def list_x11_keyboard_languages() -> Iterator[str]:
 	for line in SysCommand("localectl --no-pager list-x11-keymap-layouts", environment_vars={'SYSTEMD_COLORS': '0'}):
 		yield line.decode('UTF-8').strip()
diff --git a/archinstall/lib/menu/__init__.py b/archinstall/lib/menu/__init__.py
new file mode 100644
index 00000000..6e28c8a2
--- /dev/null
+++ b/archinstall/lib/menu/__init__.py
@@ -0,0 +1 @@
+from .menu import Menu
diff --git a/archinstall/lib/menu.py b/archinstall/lib/menu/menu.py
index 6f1c2237..65be4956 100644
--- a/archinstall/lib/menu.py
+++ b/archinstall/lib/menu/menu.py
@@ -1,4 +1,4 @@
-from .simple_menu import TerminalMenu
+from archinstall.lib.menu.simple_menu import TerminalMenu
 
 
 class Menu(TerminalMenu):
diff --git a/archinstall/lib/menu/selection_menu.py b/archinstall/lib/menu/selection_menu.py
new file mode 100644
index 00000000..17cc3347
--- /dev/null
+++ b/archinstall/lib/menu/selection_menu.py
@@ -0,0 +1,392 @@
+import sys
+
+import archinstall
+from archinstall import Menu
+
+
+class Selector:
+	def __init__(
+		self,
+		description,
+		func=None,
+		display_func=None,
+		default=None,
+		enabled=False,
+		dependencies=[],
+		dependencies_not=[]
+	):
+		"""
+		Create a new menu selection entry
+
+		:param description: Text that will be displayed as the menu entry
+		:type description: str
+
+		:param func: Function that is called when the menu entry is selected
+		:type func: Callable
+
+		:param display_func: After specifying a setting for a menu item it is displayed
+		on the right side of the item as is; with this function one can modify the entry
+		to be displayed; e.g. when specifying a password one can display **** instead
+		:type display_func: Callable
+
+		:param default: Default value for this menu entry
+		:type default: Any
+
+		:param enabled: Specify if this menu entry should be displayed
+		:type enabled: bool
+
+		:param dependencies: Specify dependencies for this menu entry; if the dependencies
+		are not set yet, then this item is not displayed; e.g. disk_layout depends on selectiong
+		harddrive(s) first
+		:type dependencies: list
+
+		:param dependencies_not: These are the exclusive options; the menu item will only be
+		displayed if non of the entries in the list have been specified
+		:type dependencies_not: list
+		"""
+
+		self._description = description
+		self.func = func
+		self._display_func = display_func
+		self._current_selection = default
+		self.enabled = enabled
+		self.text = self.menu_text()
+		self._dependencies = dependencies
+		self._dependencies_not = dependencies_not
+
+	@property
+	def dependencies(self):
+		return self._dependencies
+
+	@property
+	def dependencies_not(self):
+		return self._dependencies_not
+
+	def set_enabled(self):
+		self.enabled = True
+
+	def update_description(self, description):
+		self._description = description
+		self.text = self.menu_text()
+
+	def menu_text(self):
+		current = ''
+
+		if self._display_func:
+			current = self._display_func(self._current_selection)
+		else:
+			if self._current_selection is not None:
+				current = str(self._current_selection)
+
+		if current:
+			padding = 35 - len(self._description)
+			current = ' ' * padding + f'SET: {current}'
+
+		return f'{self._description} {current}'
+
+	def set_current_selection(self, current):
+		self._current_selection = current
+		self.text = self.menu_text()
+
+	def has_selection(self):
+		if self._current_selection is None:
+			return False
+		return True
+
+	def is_empty(self):
+		if self._current_selection is None:
+			return True
+		elif isinstance(self._current_selection, (str, list, dict)) and len(self._current_selection) == 0:
+			return True
+
+		return False
+
+
+class GlobalMenu:
+	def __init__(self):
+		self._menu_options = {}
+		self._setup_selection_menu_options()
+
+	def _setup_selection_menu_options(self):
+		self._menu_options['keyboard-layout'] = \
+			Selector('Select keyboard layout', lambda: archinstall.select_language('us'), default='us')
+		self._menu_options['mirror-region'] = \
+			Selector(
+				'Select mirror region',
+				lambda: archinstall.select_mirror_regions(),
+				display_func=lambda x: list(x.keys()) if x else '[]',
+				default={})
+		self._menu_options['sys-language'] = \
+			Selector('Select locale language', lambda: archinstall.select_locale_lang('en_US'), default='en_US')
+		self._menu_options['sys-encoding'] = \
+			Selector('Select locale encoding', lambda: archinstall.select_locale_enc('utf-8'), default='utf-8')
+		self._menu_options['harddrives'] = \
+			Selector(
+				'Select harddrives',
+				lambda: self._select_harddrives())
+		self._menu_options['disk_layouts'] = \
+			Selector(
+				'Select disk layout',
+				lambda: archinstall.select_disk_layout(
+					archinstall.arguments['harddrives'],
+					archinstall.arguments.get('advanced', False)
+				),
+				dependencies=['harddrives'])
+		self._menu_options['!encryption-password'] = \
+			Selector(
+				'Set encryption password',
+				lambda: archinstall.get_password(prompt='Enter disk encryption password (leave blank for no encryption): '),
+				display_func=lambda x: self._secret(x) if x else 'None',
+				dependencies=['harddrives'])
+		self._menu_options['swap'] = \
+			Selector(
+				'Use swap',
+				lambda: archinstall.ask_for_swap(),
+				default=True)
+		self._menu_options['bootloader'] = \
+			Selector(
+				'Select bootloader',
+				lambda: archinstall.ask_for_bootloader(archinstall.arguments.get('advanced', False)),)
+		self._menu_options['hostname'] = \
+			Selector('Specify hostname', lambda: archinstall.ask_hostname())
+		self._menu_options['!root-password'] = \
+			Selector(
+				'Set root password',
+				lambda: self._set_root_password(),
+				display_func=lambda x: self._secret(x) if x else 'None')
+		self._menu_options['!superusers'] = \
+			Selector(
+				'Specify superuser account',
+				lambda: self._create_superuser_account(),
+				dependencies_not=['!root-password'],
+				display_func=lambda x: list(x.keys()) if x else '')
+		self._menu_options['!users'] = \
+			Selector(
+				'Specify user account',
+				lambda: self._create_user_account(),
+				default={},
+				display_func=lambda x: list(x.keys()) if x else '[]')
+		self._menu_options['profile'] = \
+			Selector(
+				'Specify profile',
+				lambda: self._select_profile(),
+				display_func=lambda x: x if x else 'None')
+		self._menu_options['audio'] = \
+			Selector(
+				'Select audio',
+				lambda: archinstall.ask_for_audio_selection(archinstall.is_desktop_profile(archinstall.arguments.get('profile', None))))
+		self._menu_options['kernels'] = \
+			Selector(
+				'Select kernels',
+				lambda: archinstall.select_kernel(),
+				default='linux')
+		self._menu_options['packages'] = \
+			Selector(
+				'Additional packages to install',
+				lambda: archinstall.ask_additional_packages_to_install(archinstall.arguments.get('packages', None)),
+				default=[])
+		self._menu_options['nic'] = \
+			Selector(
+				'Configure network',
+				lambda: archinstall.ask_to_configure_network(),
+				display_func=lambda x: x if x else 'Not configured, unavailable unless setup manually',
+				default={})
+		self._menu_options['timezone'] = \
+			Selector('Select timezone', lambda: archinstall.ask_timezone())
+		self._menu_options['ntp'] = \
+			Selector(
+				'Set automatic time sync (NTP)',
+				lambda: archinstall.ask_ntp(),
+				default=True)
+		self._menu_options['install'] = \
+			Selector(
+				self._install_text(),
+				enabled=True)
+		self._menu_options['abort'] = Selector('Abort', enabled=True)
+
+	def enable(self, selector_name, omit_if_set=False):
+		arg = archinstall.arguments.get(selector_name, None)
+
+		# don't display the menu option if it was defined already
+		if arg is not None and omit_if_set:
+			return
+
+		if self._menu_options.get(selector_name, None):
+			self._menu_options[selector_name].set_enabled()
+			if arg is not None:
+				self._menu_options[selector_name].set_current_selection(arg)
+		else:
+			print(f'No selector found: {selector_name}')
+			sys.exit(1)
+
+	def run(self):
+		while True:
+			# # Before continuing, set the preferred keyboard layout/language in the current terminal.
+			# # This will just help the user with the next following questions.
+			self._set_kb_language()
+
+			enabled_menus = self._menus_to_enable()
+			menu_text = [m.text for m in enabled_menus.values()]
+			selection = Menu('Set/Modify the below options', menu_text, sort=False).run()
+			if selection:
+				selection = selection.strip()
+				if 'Abort' in selection:
+					exit(0)
+				elif 'Install' in selection:
+					if self._missing_configs() == 0:
+						self._post_processing()
+						break
+				else:
+					self._process_selection(selection)
+
+	def _process_selection(self, selection):
+		# find the selected option in our option list
+		option = [[k, v] for k, v in self._menu_options.items() if v.text.strip() == selection]
+
+		if len(option) != 1:
+			raise ValueError(f'Selection not found: {selection}')
+
+		selector_name = option[0][0]
+		selector = option[0][1]
+		result = selector.func()
+		self._menu_options[selector_name].set_current_selection(result)
+		archinstall.arguments[selector_name] = result
+
+		self._update_install()
+
+	def _update_install(self):
+		text = self._install_text()
+		self._menu_options.get('install').update_description(text)
+
+	def _post_processing(self):
+		if archinstall.arguments.get('harddrives', None) and archinstall.arguments.get('!encryption-password', None):
+			# If no partitions was marked as encrypted, but a password was supplied and we have some disks to format..
+			# Then we need to identify which partitions to encrypt. This will default to / (root).
+			if len(list(archinstall.encrypted_partitions(archinstall.storage['disk_layouts']))) == 0:
+				archinstall.storage['disk_layouts'] = archinstall.select_encrypted_partitions(
+					archinstall.storage['disk_layouts'], archinstall.arguments['!encryption-password'])
+
+	def _install_text(self):
+		missing = self._missing_configs()
+		if missing > 0:
+			return f'Install ({missing} config(s) missing)'
+		return 'Install'
+
+	def _missing_configs(self):
+		def check(s):
+			return self._menu_options.get(s).has_selection()
+
+		missing = 0
+		if not check('bootloader'):
+			missing += 1
+		if not check('hostname'):
+			missing += 1
+		if not check('audio'):
+			missing += 1
+		if not check('timezone'):
+			missing += 1
+		if not check('!root-password') and not check('!superusers'):
+			missing += 1
+		if not check('harddrives'):
+			missing += 1
+		if check('harddrives'):
+			if not self._menu_options.get('harddrives').is_empty() and not check('disk_layouts'):
+				missing += 1
+
+		return missing
+
+	def _set_root_password(self):
+		prompt = 'Enter root password (leave blank to disable root & create superuser): '
+		password = archinstall.get_password(prompt=prompt)
+
+		if password is not None:
+			self._menu_options.get('!superusers').set_current_selection(None)
+			archinstall.arguments['!users'] = {}
+			archinstall.arguments['!superusers'] = {}
+
+		return password
+
+	def _select_harddrives(self):
+		old_haddrives = archinstall.arguments.get('harddrives')
+		harddrives = archinstall.select_harddrives()
+
+		# in case the harddrives got changed we have to reset the disk layout as well
+		if old_haddrives != harddrives:
+			self._menu_options.get('disk_layouts').set_current_selection(None)
+			archinstall.arguments['disk_layouts'] = {}
+
+		if not harddrives:
+			prompt = 'You decided to skip harddrive selection\n'
+			prompt += f"and will use whatever drive-setup is mounted at {archinstall.storage['MOUNT_POINT']} (experimental)\n"
+			prompt += "WARNING: Archinstall won't check the suitability of this setup\n"
+
+			prompt += 'Do you wish to continue?'
+			choice = Menu(prompt, ['yes', 'no'], default_option='yes').run()
+
+			if choice == 'no':
+				return self._select_harddrives()
+
+		return harddrives
+
+	def _secret(self, x):
+		return '*' * len(x)
+
+	def _select_profile(self):
+		profile = archinstall.select_profile()
+
+		# Check the potentially selected profiles preparations to get early checks if some additional questions are needed.
+		if profile and profile.has_prep_function():
+			namespace = f'{profile.namespace}.py'
+			with profile.load_instructions(namespace=namespace) as imported:
+				if not imported._prep_function():
+					archinstall.log(' * Profile\'s preparation requirements was not fulfilled.', fg='red')
+					exit(1)
+
+		return profile
+
+	def _create_superuser_account(self):
+		superuser = archinstall.ask_for_superuser_account('Create a required super-user with sudo privileges: ', forced=True)
+		return superuser
+
+	def _create_user_account(self):
+		users, superusers = archinstall.ask_for_additional_users('Enter a username to create an additional user: ')
+		if not archinstall.arguments.get('!superusers', None):
+			archinstall.arguments['!superusers'] = superusers
+		else:
+			archinstall.arguments['!superusers'] = {**archinstall.arguments['!superusers'], **superusers}
+
+		return users
+
+	def _set_kb_language(self):
+		# Before continuing, set the preferred keyboard layout/language in the current terminal.
+		# This will just help the user with the next following questions.
+		if archinstall.arguments.get('keyboard-layout', None) and len(archinstall.arguments['keyboard-layout']):
+			archinstall.set_keyboard_language(archinstall.arguments['keyboard-layout'])
+
+	def _verify_selection_enabled(self, selection_name):
+		if selection := self._menu_options.get(selection_name, None):
+			if not selection.enabled:
+				return False
+
+			if len(selection.dependencies) > 0:
+				for d in selection.dependencies:
+					if not self._verify_selection_enabled(d) or self._menu_options.get(d).is_empty():
+						return False
+
+			if len(selection.dependencies_not) > 0:
+				for d in selection.dependencies_not:
+					if not self._menu_options.get(d).is_empty():
+						return False
+
+			return True
+
+		raise ValueError(f'No selection found: {selection_name}')
+
+	def _menus_to_enable(self):
+		enabled_menus = {}
+
+		for name, selection in self._menu_options.items():
+			if self._verify_selection_enabled(name):
+				enabled_menus[name] = selection
+
+		return enabled_menus
diff --git a/archinstall/lib/simple_menu.py b/archinstall/lib/menu/simple_menu.py
index a9d6d7ec..a9d6d7ec 100644
--- a/archinstall/lib/simple_menu.py
+++ b/archinstall/lib/menu/simple_menu.py
diff --git a/archinstall/lib/user_interaction.py b/archinstall/lib/user_interaction.py
index df53ce49..11ce4072 100644
--- a/archinstall/lib/user_interaction.py
+++ b/archinstall/lib/user_interaction.py
@@ -17,7 +17,7 @@ if TYPE_CHECKING:
 from .disk import BlockDevice, suggest_single_disk_layout, suggest_multi_disk_layout, valid_parted_position, all_disks
 from .exceptions import RequirementError, UserError, DiskError
 from .hardware import AVAILABLE_GFX_DRIVERS, has_uefi, has_amd_graphics, has_intel_graphics, has_nvidia_graphics
-from .locale_helpers import list_keyboard_languages, list_timezones
+from .locale_helpers import list_keyboard_languages, list_timezones, list_locales
 from .networking import list_interfaces
 from .menu import Menu
 from .output import log
@@ -27,7 +27,7 @@ from .mirrors import list_mirrors
 
 # TODO: Some inconsistencies between the selection processes.
 #       Some return the keys from the options, some the values?
-from .. import fs_types
+from .. import fs_types, validate_package_list
 
 # TODO: These can be removed after the move to simple_menu.py
 def get_terminal_height() -> int:
@@ -264,8 +264,21 @@ class MiniCurses:
 			return response
 
 
-def ask_for_swap(prompt :str = 'Would you like to use swap on zram? (Y/n): ', forced :bool = False) -> bool:
-	return True if input(prompt).strip(' ').lower() not in ('n', 'no') else False
+def ask_for_swap(prompt='Would you like to use swap on zram?', forced=False):
+	choice = Menu(prompt, ['yes', 'no'], default_option='yes').run()
+	return False if choice == 'no' else True
+
+
+def ask_ntp():
+	prompt = 'Would you like to use automatic time synchronization (NTP) with the default time servers?'
+	prompt += 'Hardware time and other post-configuration steps might be required in order for NTP to work. For more information, please check the Arch wiki'
+	choice = Menu(prompt, ['yes', 'no'], skip=False, default_option='yes').run()
+	return False if choice == 'no' else True
+
+
+def ask_hostname():
+	hostname = input('Desired hostname for the installation: ').strip(' ')
+	return hostname
 
 
 def ask_for_superuser_account(prompt :str = 'Username for required superuser with sudo privileges: ', forced :bool = False) -> Dict[str, Dict[str, str]]:
@@ -324,8 +337,8 @@ def ask_for_bootloader(advanced_options :bool = False) -> str:
 	bootloader = "systemd-bootctl" if has_uefi() else "grub-install"
 	if has_uefi():
 		if not advanced_options:
-			bootloader_choice = input("Would you like to use GRUB as a bootloader instead of systemd-boot? [y/N] ").lower()
-			if bootloader_choice == "y":
+			bootloader_choice = Menu('Would you like to use GRUB as a bootloader instead of systemd-boot?', ['yes', 'no'], default_option='no').run()
+			if bootloader_choice == "yes":
 				bootloader = "grub-install"
 		else:
 			# We use the common names for the bootloader as the selection, and map it back to the expected values.
@@ -345,10 +358,42 @@ def ask_for_bootloader(advanced_options :bool = False) -> str:
 def ask_for_audio_selection(desktop :bool = True) -> str:
 	audio = 'pipewire' if desktop else 'none'
 	choices = ['pipewire', 'pulseaudio'] if desktop else ['pipewire', 'pulseaudio', 'none']
-	selected_audio = Menu(f'Choose an audio server or leave blank to use "{audio}"', choices, default_option=audio).run()
+	selected_audio = Menu(
+		f'Choose an audio server',
+		choices,
+		default_option=audio,
+		skip=False
+	).run()
 	return selected_audio
 
 
+# TODO: Remove? Moved?
+def ask_additional_packages_to_install(packages :List[str] = None) -> List[str]:
+	# Additional packages (with some light weight error handling for invalid package names)
+	print(
+		"Only packages such as base, base-devel, linux, linux-firmware, efibootmgr and optional profile packages are installed.")
+	print("If you desire a web browser, such as firefox or chromium, you may specify it in the following prompt.")
+	while True:
+		if packages is None:
+			packages = [p for p in input(
+				'Write additional packages to install (space separated, leave blank to skip): '
+			).split(' ') if len(p)]
+
+		if len(packages):
+			# Verify packages that were given
+			try:
+				log("Verifying that additional packages exist (this might take a few seconds)")
+				validate_package_list(packages)
+				break
+			except RequirementError as e:
+				log(e, fg='red')
+		else:
+			# no additional packages were selected, which we'll allow
+			break
+
+	return packages
+
+
 def ask_to_configure_network() -> Dict[str, Any]:
 	# Optionally configure one network interface.
 	# while 1:
@@ -750,7 +795,7 @@ def select_profile() -> Optional[str]:
 	return None
 
 
-def select_language() -> str:
+def select_language(default_value :str) -> str:
 	"""
 	Asks the user to select a language
 	Usually this is combined with :ref:`archinstall.list_keyboard_languages`.
@@ -764,7 +809,7 @@ def select_language() -> str:
 	# allows for searching anyways
 	sorted_kb_lang = sorted(sorted(list(kb_lang)), key=len)
 
-	selected_lang = Menu('Select Keyboard layout', sorted_kb_lang, default_option='us', sort=False).run()
+	selected_lang = Menu('Select Keyboard layout', sorted_kb_lang, default_option=default_value, sort=False).run()
 	return selected_lang
 
 
@@ -809,7 +854,7 @@ def select_harddrives() -> Optional[str]:
 	if selected_harddrive and len(selected_harddrive) > 0:
 		return [options[i] for i in selected_harddrive]
 
-	return None
+	return []
 
 
 def select_driver(options :Dict[str, Any] = AVAILABLE_GFX_DRIVERS) -> str:
@@ -866,3 +911,31 @@ def select_kernel() -> List[str]:
 	).run()
 
 	return selected_kernels
+
+
+def select_locale_lang(default):
+	locales = list_locales()
+	locale_lang = set([locale.split()[0] for locale in locales])
+
+	selected_locale = Menu(
+		f'Choose which locale language to use',
+		locale_lang,
+		sort=True,
+		default_option=default
+	).run()
+
+	return selected_locale
+
+
+def select_locale_enc(default):
+	locales = list_locales()
+	locale_enc = set([locale.split()[1] for locale in locales])
+
+	selected_locale = Menu(
+		f'Choose which locale encoding to use',
+		locale_enc,
+		sort=True,
+		default_option=default
+	).run()
+
+	return selected_locale