1 files changed, 179 insertions, 0 deletions

diff --git a/archinstall/lib/disk/encryption_menu.py b/archinstall/lib/disk/encryption_menu.py
new file mode 100644
index 00000000..285270fb
--- /dev/null
+++ b/archinstall/lib/disk/encryption_menu.py
@@ -0,0 +1,179 @@
+from pathlib import Path
+from typing import Dict, Optional, Any, TYPE_CHECKING, List
+
+from ..disk import (
+	DeviceModification,
+	PartitionModification,
+	DiskEncryption,
+	EncryptionType
+)
+from ..menu import (
+	Selector,
+	AbstractSubMenu,
+	MenuSelectionType,
+	TableMenu
+)
+from ..user_interaction.utils import get_password
+from ..menu import Menu
+from ..general import secret
+from .fido import Fido2Device, Fido2
+from ..output import FormattedOutput
+
+if TYPE_CHECKING:
+	_: Any
+
+
+class DiskEncryptionMenu(AbstractSubMenu):
+	def __init__(
+		self,
+		mods: List[DeviceModification],
+		data_store: Dict[str, Any],
+		preset: Optional[DiskEncryption] = None
+	):
+		if preset:
+			self._preset = preset
+		else:
+			self._preset = DiskEncryption()
+
+		self._modifications = mods
+		super().__init__(data_store=data_store)
+
+	def setup_selection_menu_options(self):
+		self._menu_options['encryption_password'] = \
+			Selector(
+				_('Encryption password'),
+				lambda x: select_encrypted_password(),
+				display_func=lambda x: secret(x) if x else '',
+				default=self._preset.encryption_password,
+				enabled=True
+			)
+		self._menu_options['encryption_type'] = \
+			Selector(
+				_('Encryption type'),
+				func=lambda preset: select_encryption_type(preset),
+				display_func=lambda x: EncryptionType.type_to_text(x) if x else None,
+				dependencies=['encryption_password'],
+				default=self._preset.encryption_type,
+				enabled=True
+			)
+		self._menu_options['partitions'] = \
+			Selector(
+				_('Partitions'),
+				func=lambda preset: select_partitions_to_encrypt(self._modifications.device_modifications, preset),
+				display_func=lambda x: f'{len(x)} {_("Partitions")}' if x else None,
+				dependencies=['encryption_password'],
+				default=self._preset.partitions,
+				preview_func=self._prev_disk_layouts,
+				enabled=True
+			)
+		self._menu_options['HSM'] = \
+			Selector(
+				description=_('Use HSM to unlock encrypted drive'),
+				func=lambda preset: select_hsm(preset),
+				display_func=lambda x: self._display_hsm(x),
+				dependencies=['encryption_password'],
+				default=self._preset.hsm_device,
+				enabled=True
+			)
+
+	def run(self, allow_reset: bool = True) -> Optional[DiskEncryption]:
+		super().run(allow_reset=allow_reset)
+
+		if self._data_store.get('encryption_password', None):
+			return DiskEncryption(
+				encryption_password=self._data_store.get('encryption_password', None),
+				encryption_type=self._data_store['encryption_type'],
+				partitions=self._data_store.get('partitions', None),
+				hsm_device=self._data_store.get('HSM', None)
+			)
+
+		return None
+
+	def _display_hsm(self, device: Optional[Fido2Device]) -> Optional[str]:
+		if device:
+			return device.manufacturer
+
+		if not Fido2.get_fido2_devices():
+			return str(_('No HSM devices available'))
+		return None
+
+	def _prev_disk_layouts(self) -> Optional[str]:
+		partitions: Optional[List[PartitionModification]] = self._menu_options['partitions'].current_selection
+		if partitions:
+			output = str(_('Partitions to be encrypted')) + '\n'
+			output += FormattedOutput.as_table(partitions)
+			return output.rstrip()
+
+		return None
+
+
+def select_encryption_type(preset: EncryptionType) -> Optional[EncryptionType]:
+	title = str(_('Select disk encryption option'))
+	options = [
+		EncryptionType.type_to_text(EncryptionType.Partition)
+	]
+
+	preset_value = EncryptionType.type_to_text(preset)
+	choice = Menu(title, options, preset_values=preset_value).run()
+
+	match choice.type_:
+		case MenuSelectionType.Reset: return None
+		case MenuSelectionType.Skip: return preset
+		case MenuSelectionType.Selection: return EncryptionType.text_to_type(choice.value)  # type: ignore
+
+
+def select_encrypted_password() -> Optional[str]:
+	if passwd := get_password(prompt=str(_('Enter disk encryption password (leave blank for no encryption): '))):
+		return passwd
+	return None
+
+
+def select_hsm(preset: Optional[Fido2Device] = None) -> Optional[Fido2Device]:
+	title = _('Select a FIDO2 device to use for HSM')
+	fido_devices = Fido2.get_fido2_devices()
+
+	if fido_devices:
+		choice = TableMenu(title, data=fido_devices).run()
+		match choice.type_:
+			case MenuSelectionType.Reset:
+				return None
+			case MenuSelectionType.Skip:
+				return preset
+			case MenuSelectionType.Selection:
+				return choice.value  # type: ignore
+
+	return None
+
+
+def select_partitions_to_encrypt(
+	modification: List[DeviceModification],
+	preset: List[PartitionModification]
+) -> List[PartitionModification]:
+	partitions: List[PartitionModification] = []
+
+	# do not allow encrypting the boot partition
+	for mod in modification:
+		partitions += list(filter(lambda x: x.mountpoint != Path('/boot'), mod.partitions))
+
+	# do not allow encrypting existing partitions that are not marked as wipe
+	avail_partitions = list(filter(lambda x: not x.exists(), partitions))
+
+	if avail_partitions:
+		title = str(_('Select which partitions to encrypt'))
+		partition_table = FormattedOutput.as_table(avail_partitions)
+
+		choice = TableMenu(
+			title,
+			table_data=(avail_partitions, partition_table),
+			preset=preset,
+			multi=True
+		).run()
+
+		match choice.type_:
+			case MenuSelectionType.Reset:
+				return []
+			case MenuSelectionType.Skip:
+				return preset
+			case MenuSelectionType.Selection:
+				return choice.multi_value
+	return []