New encryption menu (#1520)

* New encryption menu

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

5 files changed, 226 insertions, 112 deletions

diff --git a/archinstall/lib/menu/selection_menu.py b/archinstall/lib/menu/abstract_menu.py
index 8a08812c..61466e3e 100644
--- a/archinstall/lib/menu/selection_menu.py
+++ b/archinstall/lib/menu/abstract_menu.py
@@ -2,15 +2,12 @@ from __future__ import annotations
 
 import logging
 import sys
-import pathlib
 from typing import Callable, Any, List, Iterator, Tuple, Optional, Dict, TYPE_CHECKING
 
 from .menu import Menu, MenuSelectionType
 from ..locale_helpers import set_keyboard_language
 from ..output import log
 from ..translationhandler import TranslationHandler, Language
-from ..hsm.fido import get_fido2_devices
-
 from ..user_interaction.general_conf import select_archinstall_language
 
 if TYPE_CHECKING:
@@ -167,8 +164,9 @@ class Selector:
 		if status and not self.is_enabled():
 			self.set_enabled(True)
 
-class GeneralMenu:
-	def __init__(self, data_store :dict = None, auto_cursor=False, preview_size :float = 0.2):
+
+class AbstractMenu:
+	def __init__(self, data_store: Dict[str, Any] = None, auto_cursor=False, preview_size :float = 0.2):
 		"""
 		Create a new selection menu.
 
@@ -196,7 +194,7 @@ class GeneralMenu:
 	def last_choice(self):
 		return self._last_choice
 
-	def __enter__(self, *args :Any, **kwargs :Any) -> GeneralMenu:
+	def __enter__(self, *args :Any, **kwargs :Any) -> AbstractMenu:
 		self.is_context_mgr = True
 		return self
 
@@ -209,9 +207,9 @@ class GeneralMenu:
 			raise args[1]
 
 		for key in self._menu_options:
-			sel = self._menu_options[key]
+			selector = self._menu_options[key]
 			if key and key not in self._data_store:
-				self._data_store[key] = sel._current_selection
+				self._data_store[key] = selector.current_selection
 
 		self.exit_callback()
 
@@ -286,7 +284,7 @@ class GeneralMenu:
 		selector = option[0][1]
 		return config_name, selector
 
-	def run(self):
+	def run(self, allow_reset: bool = False):
 		""" Calls the Menu framework"""
 		# we synch all the options just in case
 		for item in self.list_options():
@@ -304,6 +302,8 @@ class GeneralMenu:
 			padding = self._get_menu_text_padding(list(enabled_menus.values()))
 			menu_options = [m.menu_text(padding) for m in enabled_menus.values()]
 
+			warning_msg = str(_('All settings will be reset, are you sure?'))
+
 			selection = Menu(
 				_('Set/Modify the below options'),
 				menu_options,
@@ -312,33 +312,39 @@ class GeneralMenu:
 				preview_command=self._preview_display,
 				preview_size=self.preview_size,
 				skip_empty_entries=True,
-				skip=False
+				skip=False,
+				allow_reset=allow_reset,
+				allow_reset_warning_msg=warning_msg
 			).run()
 
-			if selection.type_ == MenuSelectionType.Selection:
-				value = selection.value
+			match selection.type_:
+				case MenuSelectionType.Reset:
+					self._data_store = {}
+					return
+				case MenuSelectionType.Selection:
+					value: str = selection.value  # type: ignore
 
-				if self.auto_cursor:
-					cursor_pos = menu_options.index(value) + 1  # before the strip otherwise fails
+					if self.auto_cursor:
+						cursor_pos = menu_options.index(value) + 1  # before the strip otherwise fails
 
-					# in case the new position lands on a "placeholder" we'll skip them as well
-					while True:
-						if cursor_pos >= len(menu_options):
-							cursor_pos = 0
-						if len(menu_options[cursor_pos]) > 0:
-							break
-						cursor_pos += 1
+						# in case the new position lands on a "placeholder" we'll skip them as well
+						while True:
+							if cursor_pos >= len(menu_options):
+								cursor_pos = 0
+							if len(menu_options[cursor_pos]) > 0:
+								break
+							cursor_pos += 1
 
-				value = value.strip()
+					value = value.strip()
 
-				# if this calls returns false, we exit the menu
-				# we allow for an callback for special processing on realeasing control
-				if not self._process_selection(value):
-					break
+					# if this calls returns false, we exit the menu
+					# we allow for an callback for special processing on realeasing control
+					if not self._process_selection(value):
+						break
 
 		# we get the last action key
 		actions = {str(v.description):k for k,v in self._menu_options.items()}
-		self._last_choice = actions[selection.value.strip()]
+		self._last_choice = actions[selection.value.strip()]  # type: ignore
 
 		if not self.is_context_mgr:
 			self.__exit__()
@@ -472,26 +478,16 @@ class GeneralMenu:
 		self._translation_handler.activate(language)
 		return language
 
-	def _select_hsm(self, preset :Optional[pathlib.Path] = None) -> Optional[pathlib.Path]:
-		title = _('Select which partitions to mark for formatting:')
-		title += '\n'
-
-		fido_devices = get_fido2_devices()
 
-		indexes = []
-		for index, path in enumerate(fido_devices.keys()):
-			title += f"{index}: {path} ({fido_devices[path]['manufacturer']} - {fido_devices[path]['product']})"
-			indexes.append(f"{index}|{fido_devices[path]['product']}")
+class AbstractSubMenu(AbstractMenu):
+	def __init__(self, data_store: Dict[str, Any] = None):
+		super().__init__(data_store=data_store)
 
-		title += '\n'
-
-		choice = Menu(title, indexes, multi=False).run()
-
-		match choice.type_:
-			case MenuSelectionType.Esc: return preset
-			case MenuSelectionType.Selection:
-				selection: Any = choice.value
-				index = int(selection.split('|',1)[0])
-				return pathlib.Path(list(fido_devices.keys())[index])
-
-		return None
+		self._menu_options['__separator__'] = Selector('')
+		self._menu_options['back'] = \
+			Selector(
+				_('Back'),
+				no_store=True,
+				enabled=True,
+				exec_func=lambda n, v: True,
+			)
diff --git a/archinstall/lib/menu/global_menu.py b/archinstall/lib/menu/global_menu.py
index 444ba7ee..0d348227 100644
--- a/archinstall/lib/menu/global_menu.py
+++ b/archinstall/lib/menu/global_menu.py
@@ -3,12 +3,13 @@ from __future__ import annotations
 from typing import Any, List, Optional, Union, Dict, TYPE_CHECKING
 
 import archinstall
-from ..disk import encrypted_partitions
+from ..disk.encryption import DiskEncryptionMenu
 from ..general import SysCommand, secret
 from ..hardware import has_uefi
 from ..menu import Menu
-from ..menu.selection_menu import Selector, GeneralMenu
+from ..menu.abstract_menu import Selector, AbstractMenu
 from ..models import NetworkConfiguration
+from ..models.disk_encryption import DiskEncryption, EncryptionType
 from ..models.users import User
 from ..output import FormattedOutput
 from ..profiles import is_desktop_profile, Profile
@@ -25,7 +26,6 @@ from ..user_interaction import ask_to_configure_network
 from ..user_interaction import get_password, ask_for_a_timezone, save_config
 from ..user_interaction import select_additional_repositories
 from ..user_interaction import select_disk_layout
-from ..user_interaction import select_encrypted_partitions
 from ..user_interaction import select_harddrives
 from ..user_interaction import select_kernel
 from ..user_interaction import select_language
@@ -39,7 +39,7 @@ if TYPE_CHECKING:
 	_: Any
 
 
-class GlobalMenu(GeneralMenu):
+class GlobalMenu(AbstractMenu):
 	def __init__(self,data_store):
 		self._disk_check = True
 		super().__init__(data_store=data_store, auto_cursor=True, preview_size=0.3)
@@ -91,18 +91,13 @@ class GlobalMenu(GeneralMenu):
 				preview_func=self._prev_disk_layouts,
 				display_func=lambda x: self._display_disk_layout(x),
 				dependencies=['harddrives'])
-		self._menu_options['!encryption-password'] = \
+		self._menu_options['disk_encryption'] = \
 			Selector(
-				_('Encryption password'),
-				lambda x: self._select_encrypted_password(),
-				display_func=lambda x: secret(x) if x else 'None',
-				dependencies=['harddrives'])
-		self._menu_options['HSM'] = Selector(
-			description=_('Use HSM to unlock encrypted drive'),
-			func=lambda preset: self._select_hsm(preset),
-			dependencies=['!encryption-password'],
-			default=None
-		)
+				_('Disk encryption'),
+				lambda preset: self._disk_encryption(preset),
+				preview_func=self._prev_disk_encryption,
+				display_func=lambda x: self._display_disk_encryption(x),
+				dependencies=['disk_layouts'])
 		self._menu_options['swap'] = \
 			Selector(
 				_('Swap'),
@@ -209,28 +204,6 @@ class GlobalMenu(GeneralMenu):
 	def post_callback(self,name :str = None ,result :Any = None):
 		self._update_install_text(name, result)
 
-	def exit_callback(self):
-		if self._data_store.get('harddrives', None) and self._data_store.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(encrypted_partitions(storage['arguments'].get('disk_layouts', [])))) == 0:
-				for blockdevice in storage['arguments']['disk_layouts']:
-					if storage['arguments']['disk_layouts'][blockdevice].get('partitions'):
-						for partition_index in select_encrypted_partitions(
-								title=_('Select which partitions to encrypt:'),
-								partitions=storage['arguments']['disk_layouts'][blockdevice]['partitions'],
-								filter_=(lambda p: p['mountpoint'] != '/boot')
-							):
-
-							partition = storage['arguments']['disk_layouts'][blockdevice]['partitions'][partition_index]
-							partition['encrypted'] = True
-							partition['!password'] = storage['arguments']['!encryption-password']
-
-							# We make sure generate-encryption-key-file is set on additional partitions
-							# other than the root partition. Otherwise they won't unlock properly #1279
-							if partition['mountpoint'] != '/':
-								partition['generate-encryption-key-file'] = True
-
 	def _install_text(self):
 		missing = len(self._missing_configs())
 		if missing > 0:
@@ -246,6 +219,20 @@ class GlobalMenu(GeneralMenu):
 			else:
 				return str(cur_value)
 
+	def _disk_encryption(self, preset: Optional[DiskEncryption]) -> Optional[DiskEncryption]:
+		data_store: Dict[str, Any] = {}
+
+		selector = self._menu_options['disk_layouts']
+
+		if selector.has_selection():
+			layouts: Dict[str, Dict[str, Any]] = selector.current_selection
+		else:
+			# this should not happen as the encryption menu has the disk layout as dependency
+			raise ValueError('No disk layout specified')
+
+		disk_encryption = DiskEncryptionMenu(data_store, preset, layouts).run()
+		return disk_encryption
+
 	def _prev_network_config(self) -> Optional[str]:
 		selector = self._menu_options['nic']
 		if selector.has_selection():
@@ -283,6 +270,30 @@ class GlobalMenu(GeneralMenu):
 			return f'{total_nr} {_("Partitions")}'
 		return ''
 
+	def _prev_disk_encryption(self) -> Optional[str]:
+		selector = self._menu_options['disk_encryption']
+		if selector.has_selection():
+			encryption: DiskEncryption = selector.current_selection
+
+			enc_type = EncryptionType.type_to_text(encryption.encryption_type)
+			output = str(_('Encryption type')) + f': {enc_type}\n'
+			output += str(_('Password')) + f': {secret(encryption.encryption_password)}\n'
+
+			if encryption.partitions:
+				output += 'Partitions: {} selected'.format(len(encryption.partitions)) + '\n'
+
+			if encryption.hsm_device:
+				output += f'HSM: {encryption.hsm_device.manufacturer}'
+
+			return output
+
+		return None
+
+	def _display_disk_encryption(self, current_value: Optional[DiskEncryption]) -> str:
+		if current_value:
+			return EncryptionType.type_to_text(current_value.encryption_type)
+		return ''
+
 	def _prev_install_missing_config(self) -> Optional[str]:
 		if missing := self._missing_configs():
 			text = str(_('Missing configurations:\n'))
@@ -327,11 +338,10 @@ class GlobalMenu(GeneralMenu):
 		password = get_password(prompt=prompt)
 		return password
 
-	def _select_encrypted_password(self) -> Optional[str]:
-		if passwd := get_password(prompt=str(_('Enter disk encryption password (leave blank for no encryption): '))):
-			return passwd
-		else:
-			return None
+	# def _select_encrypted_password(self) -> Optional[str]:
+	# 	if passwd := get_password(prompt=str(_('Enter disk encryption password (leave blank for no encryption): '))):
+	# 		return passwd
+	# 	return None
 
 	def _select_ntp(self, preset :bool = True) -> bool:
 		ntp = ask_ntp(preset)
diff --git a/archinstall/lib/menu/list_manager.py b/archinstall/lib/menu/list_manager.py
index ae3a6eb5..1e09d987 100644
--- a/archinstall/lib/menu/list_manager.py
+++ b/archinstall/lib/menu/list_manager.py
@@ -104,7 +104,7 @@ class ListManager:
 		return options, header
 
 	def _run_actions_on_entry(self, entry: Any):
-		options = self.filter_options(entry,self._sub_menu_actions) + [self._cancel_action]
+		options = self.filter_options(entry, self._sub_menu_actions) + [self._cancel_action]
 		display_value = self.selected_action_display(entry)
 
 		prompt = _("Select an action for '{}'").format(display_value)
diff --git a/archinstall/lib/menu/menu.py b/archinstall/lib/menu/menu.py
index 773ff1de..09685c55 100644
--- a/archinstall/lib/menu/menu.py
+++ b/archinstall/lib/menu/menu.py
@@ -18,8 +18,8 @@ if TYPE_CHECKING:
 
 class MenuSelectionType(Enum):
 	Selection = auto()
-	Esc = auto()
-	Ctrl_c = auto()
+	Skip = auto()
+	Reset = auto()
 
 
 @dataclass
@@ -56,8 +56,8 @@ class Menu(TerminalMenu):
 		preview_size: float = 0.0,
 		preview_title: str = 'Info',
 		header :Union[List[str],str] = None,
-		raise_error_on_interrupt :bool = False,
-		raise_error_warning_msg :str = '',
+		allow_reset :bool = False,
+		allow_reset_warning_msg :str = '',
 		clear_screen: bool = True,
 		show_search_hint: bool = True,
 		cycle_cursor: bool = True,
@@ -150,17 +150,10 @@ class Menu(TerminalMenu):
 		self._skip = skip
 		self._default_option = default_option
 		self._multi = multi
-		self._raise_error_on_interrupt = raise_error_on_interrupt
-		self._raise_error_warning_msg = raise_error_warning_msg
+		self._raise_error_on_interrupt = allow_reset
+		self._raise_error_warning_msg = allow_reset_warning_msg
 		self._preview_command = preview_command
 
-		menu_title = f'\n{title}\n\n'
-
-		if header:
-			if not isinstance(header,(list,tuple)):
-				header = [header]
-			menu_title += '\n'.join(header)
-
 		action_info = ''
 		if skip:
 			action_info += str(_('ESC to skip'))
@@ -173,7 +166,15 @@ class Menu(TerminalMenu):
 			action_info += ', ' if len(action_info) > 0 else ''
 			action_info += str(_('TAB to select'))
 
-		menu_title += action_info + '\n'
+		if action_info:
+			action_info += '\n\n'
+
+		menu_title = f'\n{action_info}{title}\n'
+
+		if header:
+			if not isinstance(header,(list,tuple)):
+				header = [header]
+			menu_title += '\n' + '\n'.join(header)
 
 		if default_option:
 			# if a default value was specified we move that one
@@ -215,7 +216,7 @@ class Menu(TerminalMenu):
 		try:
 			idx = self.show()
 		except KeyboardInterrupt:
-			return MenuSelection(type_=MenuSelectionType.Ctrl_c)
+			return MenuSelection(type_=MenuSelectionType.Reset)
 
 		def check_default(elem):
 			if self._default_option is not None and f'{self._default_option} {self._default_str}' in elem:
@@ -234,7 +235,7 @@ class Menu(TerminalMenu):
 				result = check_default(self._menu_options[idx])
 				return MenuSelection(type_=MenuSelectionType.Selection, value=result)
 		else:
-			return MenuSelection(type_=MenuSelectionType.Esc)
+			return MenuSelection(type_=MenuSelectionType.Skip)
 
 	def _preview_wrapper(self, preview_command: Optional[Callable], current_selection: str) -> Optional[str]:
 		if preview_command:
@@ -246,15 +247,15 @@ class Menu(TerminalMenu):
 	def run(self) -> MenuSelection:
 		ret = self._show()
 
-		if ret.type_ == MenuSelectionType.Ctrl_c:
+		if ret.type_ == MenuSelectionType.Reset:
 			if self._raise_error_on_interrupt and len(self._raise_error_warning_msg) > 0:
 				response = Menu(self._raise_error_warning_msg, Menu.yes_no(), skip=False).run()
 				if response.value == Menu.no():
 					return self.run()
-
-		if ret.type_ is not MenuSelectionType.Selection and not self._skip:
-			system('clear')
-			return self.run()
+		elif ret.type_ is MenuSelectionType.Skip:
+			if not self._skip:
+				system('clear')
+				return self.run()
 
 		return ret
 
diff --git a/archinstall/lib/menu/table_selection_menu.py b/archinstall/lib/menu/table_selection_menu.py
new file mode 100644
index 00000000..09cd6ee2
--- /dev/null
+++ b/archinstall/lib/menu/table_selection_menu.py
@@ -0,0 +1,107 @@
+from typing import Any, Tuple, List, Dict, Optional
+
+from .menu import MenuSelectionType, MenuSelection
+from ..output import FormattedOutput
+from ..menu import Menu
+
+
+class TableMenu(Menu):
+	def __init__(
+		self,
+		title: str,
+		data: List[Any] = [],
+		table_data: Optional[Tuple[List[Any], str]] = None,
+		custom_menu_options: List[str] = [],
+		default: Any = None,
+		multi: bool = False
+	):
+		"""
+		param title: Text that will be displayed above the menu
+		:type title: str
+
+		param data: List of objects that will be displayed as rows
+		:type data: List
+
+		param table_data: Tuple containing a list of objects and the corresponding
+		Table representation of the data as string; this can be used in case the table
+		has to be crafted in a more sophisticated manner
+		:type table_data: Optional[Tuple[List[Any], str]]
+
+		param custom_options: List of custom options that will be displayed under the table
+		:type custom_menu_options: List
+		"""
+		if not data and not table_data:
+			raise ValueError('Either "data" or "table_data" must be provided')
+
+		self._custom_options = custom_menu_options
+		self._multi = multi
+
+		if multi:
+			header_padding = 7
+		else:
+			header_padding = 2
+
+		if len(data):
+			table_text = FormattedOutput.as_table(data)
+			rows = table_text.split('\n')
+			table = self._create_table(data, rows, header_padding=header_padding)
+		elif table_data is not None:
+			# we assume the table to be
+			# h1  |   h2
+			# -----------
+			# r1  |   r2
+			data = table_data[0]
+			rows = table_data[1].split('\n')
+			table = self._create_table(data, rows, header_padding=header_padding)
+
+		self._options, header = self._prepare_selection(table)
+
+		super().__init__(
+			title,
+			self._options,
+			header=header,
+			skip_empty_entries=True,
+			show_search_hint=False,
+			allow_reset=True,
+			multi=multi,
+			default_option=default
+		)
+
+	def run(self) -> MenuSelection:
+		choice = super().run()
+
+		match choice.type_:
+			case MenuSelectionType.Selection:
+				if self._multi:
+					choice.value = [self._options[val] for val in choice.value]  # type: ignore
+				else:
+					choice.value = self._options[choice.value]  # type: ignore
+
+		return choice
+
+	def _create_table(self, data: List[Any], rows: List[str], header_padding: int = 2) -> Dict[str, Any]:
+		# these are the header rows of the table and do not map to any data obviously
+		# we're adding 2 spaces as prefix because the menu selector '> ' will be put before
+		# the selectable rows so the header has to be aligned
+		padding = ' ' * header_padding
+		display_data = {f'{padding}{rows[0]}': None, f'{padding}{rows[1]}': None}
+
+		for row, entry in zip(rows[2:], data):
+			row = row.replace('|', '\\|')
+			display_data[row] = entry
+
+		return display_data
+
+	def _prepare_selection(self, table: Dict[str, Any]) -> Tuple[Dict[str, Any], str]:
+		# header rows are mapped to None so make sure to exclude those from the selectable data
+		options = {key: val for key, val in table.items() if val is not None}
+		header = ''
+
+		if len(options) > 0:
+			table_header = [key for key, val in table.items() if val is None]
+			header = '\n'.join(table_header)
+
+		custom = {key: None for key in self._custom_options}
+		options.update(custom)
+
+		return options, header