New encryption menu (#1520)

* New encryption menu

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

4 files changed, 169 insertions, 23 deletions

diff --git a/archinstall/lib/disk/encryption.py b/archinstall/lib/disk/encryption.py
new file mode 100644
index 00000000..67f656c8
--- /dev/null
+++ b/archinstall/lib/disk/encryption.py
@@ -0,0 +1,162 @@
+from typing import Dict, Optional, Any, TYPE_CHECKING, List
+
+from ..menu.abstract_menu import Selector, AbstractSubMenu
+from ..menu.menu import MenuSelectionType
+from ..menu.table_selection_menu import TableMenu
+from ..models.disk_encryption import EncryptionType, DiskEncryption
+from ..user_interaction.partitioning_conf import current_partition_layout
+from ..user_interaction.utils import get_password
+from ..menu import Menu
+from ..general import secret
+from ..hsm.fido import Fido2Device, Fido2
+
+if TYPE_CHECKING:
+	_: Any
+
+
+class DiskEncryptionMenu(AbstractSubMenu):
+	def __init__(self, data_store: Dict[str, Any], preset: Optional[DiskEncryption], disk_layouts: Dict[str, Any]):
+		if preset:
+			self._preset = preset
+		else:
+			self._preset = DiskEncryption()
+
+		self._disk_layouts = disk_layouts
+		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._disk_layouts, 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]:
+		selector = self._menu_options['partitions']
+		if selector.has_selection():
+			partitions: List[Any] = selector.current_selection
+			output = str(_('Partitions to be encrypted')) + '\n'
+			output += current_partition_layout(partitions, with_title=False)
+			return output.rstrip()
+		return None
+
+
+def select_encryption_type(preset: EncryptionType) -> Optional[EncryptionType]:
+	title = str(_('Select disk encryption option'))
+	options = [
+		# _type_to_text(EncryptionType.FullDiskEncryption),
+		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(disk_layouts: Dict[str, Any], preset: List[Any]) -> List[Any]:
+	# 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).
+	all_partitions = []
+	for blockdevice in disk_layouts.values():
+		if partitions := blockdevice.get('partitions'):
+			partitions = [p for p in partitions if p['mountpoint'] != '/boot']
+			all_partitions += partitions
+
+	if all_partitions:
+		title = str(_('Select which partitions to encrypt'))
+		partition_table = current_partition_layout(all_partitions, with_title=False).strip()
+
+		choice = TableMenu(
+			title,
+			table_data=(all_partitions, partition_table),
+			multi=True
+		).run()
+
+		match choice.type_:
+			case MenuSelectionType.Reset:
+				return []
+			case MenuSelectionType.Skip:
+				return preset
+			case MenuSelectionType.Selection:
+				return choice.value  # type: ignore
+
+	return []
diff --git a/archinstall/lib/disk/filesystem.py b/archinstall/lib/disk/filesystem.py
index af5879aa..bdfa502a 100644
--- a/archinstall/lib/disk/filesystem.py
+++ b/archinstall/lib/disk/filesystem.py
@@ -5,6 +5,8 @@ import json
 import pathlib
 from typing import Optional, Dict, Any, TYPE_CHECKING
 # https://stackoverflow.com/a/39757388/929999
+from ..models.disk_encryption import DiskEncryption
+
 if TYPE_CHECKING:
 	from .blockdevice import BlockDevice
 	_: Any
@@ -107,33 +109,22 @@ class Filesystem:
 				continue
 
 			if partition.get('filesystem', {}).get('format', False):
-
 				# needed for backward compatibility with the introduction of the new "format_options"
 				format_options = partition.get('options',[]) + partition.get('filesystem',{}).get('format_options',[])
-				if partition.get('encrypted', False):
+				disk_encryption: DiskEncryption = storage['arguments'].get('disk_encryption')
+
+				if partition in disk_encryption.partitions:
 					if not partition['device_instance']:
 						raise DiskError(f"Internal error caused us to loose the partition. Please report this issue upstream!")
 
-					if not partition.get('!password'):
-						if not storage['arguments'].get('!encryption-password'):
-							if storage['arguments'] == 'silent':
-								raise ValueError(f"Missing encryption password for {partition['device_instance']}")
-
-							from ..user_interaction import get_password
-
-							prompt = str(_('Enter a encryption password for {}').format(partition['device_instance']))
-							storage['arguments']['!encryption-password'] = get_password(prompt)
-
-						partition['!password'] = storage['arguments']['!encryption-password']
-
 					if partition.get('mountpoint',None):
 						loopdev = f"{storage.get('ENC_IDENTIFIER', 'ai')}{pathlib.Path(partition['mountpoint']).name}loop"
 					else:
 						loopdev = f"{storage.get('ENC_IDENTIFIER', 'ai')}{pathlib.Path(partition['device_instance'].path).name}"
 
-					partition['device_instance'].encrypt(password=partition['!password'])
+					partition['device_instance'].encrypt(password=disk_encryption.encryption_password)
 					# Immediately unlock the encrypted device to format the inner volume
-					with luks2(partition['device_instance'], loopdev, partition['!password'], auto_unmount=True) as unlocked_device:
+					with luks2(partition['device_instance'], loopdev, disk_encryption.encryption_password, auto_unmount=True) as unlocked_device:
 						if not partition.get('wipe'):
 							if storage['arguments'] == 'silent':
 								raise ValueError(f"Missing fs-type to format on newly created encrypted partition {partition['device_instance']}")
diff --git a/archinstall/lib/disk/helpers.py b/archinstall/lib/disk/helpers.py
index 256f7abb..a5164b76 100644
--- a/archinstall/lib/disk/helpers.py
+++ b/archinstall/lib/disk/helpers.py
@@ -469,12 +469,6 @@ def disk_layouts() -> Optional[Dict[str, Any]]:
 		return None
 
 
-def encrypted_partitions(blockdevices :Dict[str, Any]) -> bool:
-	for blockdevice in blockdevices.values():
-		for partition in blockdevice.get('partitions', []):
-			if partition.get('encrypted', False):
-				yield partition
-
 def find_partition_by_mountpoint(block_devices :List[BlockDevice], relative_mountpoint :str) -> Partition:
 	for device in block_devices:
 		for partition in block_devices[device]['partitions']:
diff --git a/archinstall/lib/disk/partition.py b/archinstall/lib/disk/partition.py
index 04d33453..9febf102 100644
--- a/archinstall/lib/disk/partition.py
+++ b/archinstall/lib/disk/partition.py
@@ -91,7 +91,6 @@ class Partition:
 		self._path = path
 		self._part_id = part_id
 		self._target_mountpoint = mountpoint
-		self._encrypted = None
 		self._encrypted = encrypted
 		self._wipe = False
 		self._type = 'primary'