Save encryption configuration (#1672)

* Save encryption configuration

* Fix deserialization problem

* Added .part_uuid to MapperDev

---------

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

9 files changed, 108 insertions, 22 deletions

diff --git a/archinstall/__init__.py b/archinstall/__init__.py
index b22b4663..9664c507 100644
--- a/archinstall/__init__.py
+++ b/archinstall/__init__.py
@@ -240,6 +240,14 @@ def load_config():
 		superusers = arguments.get('!superusers', None)
 		arguments['!users'] = User.parse_arguments(users, superusers)
 
+	if arguments.get('disk_encryption', None) is not None and arguments.get('disk_layouts', None) is not None:
+		password = arguments.get('encryption_password', '')
+		arguments['disk_encryption'] = DiskEncryption.parse_arg(
+			arguments['disk_layouts'],
+			arguments['disk_encryption'],
+			password
+		)
+
 
 def post_process_arguments(arguments):
 	storage['arguments'] = arguments
diff --git a/archinstall/lib/configuration.py b/archinstall/lib/configuration.py
index ad537b21..e0b5a143 100644
--- a/archinstall/lib/configuration.py
+++ b/archinstall/lib/configuration.py
@@ -44,7 +44,7 @@ class ConfigurationOutput:
 		self._disk_layout_file = "user_disk_layout.json"
 
 		self._sensitive = ['!users']
-		self._ignore = ['abort', 'install', 'config', 'creds', 'dry_run', 'disk_encryption']
+		self._ignore = ['abort', 'install', 'config', 'creds', 'dry_run']
 
 		self._process_config()
 
@@ -71,6 +71,9 @@ class ConfigurationOutput:
 			else:
 				self._user_config[key] = self._config[key]
 
+				if key == 'disk_encryption':  # special handling for encryption password
+					self._user_credentials['encryption_password'] = self._config[key].encryption_password
+
 	def user_config_to_json(self) -> str:
 		return json.dumps({
 			'config_version': storage['__version__'],  # Tells us what version was used to generate the config
diff --git a/archinstall/lib/disk/encryption.py b/archinstall/lib/disk/encryption.py
index 67f656c8..c7496bfa 100644
--- a/archinstall/lib/disk/encryption.py
+++ b/archinstall/lib/disk/encryption.py
@@ -46,7 +46,7 @@ class DiskEncryptionMenu(AbstractSubMenu):
 			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,
+				display_func=lambda x: f'{sum([len(y) for y in x.values()])} {_("Partitions")}' if x else None,
 				dependencies=['encryption_password'],
 				default=self._preset.partitions,
 				preview_func=self._prev_disk_layouts,
@@ -86,9 +86,14 @@ class DiskEncryptionMenu(AbstractSubMenu):
 	def _prev_disk_layouts(self) -> Optional[str]:
 		selector = self._menu_options['partitions']
 		if selector.has_selection():
-			partitions: List[Any] = selector.current_selection
+			partitions: Dict[str, Any] = selector.current_selection
+
+			all_partitions = []
+			for parts in partitions.values():
+				all_partitions += parts
+
 			output = str(_('Partitions to be encrypted')) + '\n'
-			output += current_partition_layout(partitions, with_title=False)
+			output += current_partition_layout(all_partitions, with_title=False)
 			return output.rstrip()
 		return None
 
@@ -132,7 +137,7 @@ def select_hsm(preset: Optional[Fido2Device] = None) -> Optional[Fido2Device]:
 	return None
 
 
-def select_partitions_to_encrypt(disk_layouts: Dict[str, Any], preset: List[Any]) -> List[Any]:
+def select_partitions_to_encrypt(disk_layouts: Dict[str, Any], preset: Dict[str, Any]) -> Dict[str, 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 = []
@@ -153,10 +158,17 @@ def select_partitions_to_encrypt(disk_layouts: Dict[str, Any], preset: List[Any]
 
 		match choice.type_:
 			case MenuSelectionType.Reset:
-				return []
+				return {}
 			case MenuSelectionType.Skip:
 				return preset
 			case MenuSelectionType.Selection:
-				return choice.value  # type: ignore
+				selections: List[Any] = choice.value  # type: ignore
+				partitions = {}
+
+				for path, device in disk_layouts.items():
+					for part in selections:
+						if part in device.get('partitions', []):
+							partitions.setdefault(path, []).append(part)
 
-	return []
+				return partitions
+	return {}
diff --git a/archinstall/lib/disk/filesystem.py b/archinstall/lib/disk/filesystem.py
index 1e722ce5..1083df53 100644
--- a/archinstall/lib/disk/filesystem.py
+++ b/archinstall/lib/disk/filesystem.py
@@ -113,7 +113,7 @@ class Filesystem:
 				format_options = partition.get('options',[]) + partition.get('filesystem',{}).get('format_options',[])
 				disk_encryption: DiskEncryption = storage['arguments'].get('disk_encryption')
 
-				if disk_encryption and partition in disk_encryption.partitions:
+				if disk_encryption and partition in disk_encryption.all_partitions:
 					if not partition['device_instance']:
 						raise DiskError(f"Internal error caused us to loose the partition. Please report this issue upstream!")
 
diff --git a/archinstall/lib/hsm/fido.py b/archinstall/lib/hsm/fido.py
index 758a2548..1c226322 100644
--- a/archinstall/lib/hsm/fido.py
+++ b/archinstall/lib/hsm/fido.py
@@ -1,9 +1,11 @@
+from __future__ import annotations
+
 import getpass
 import logging
 
 from dataclasses import dataclass
 from pathlib import Path
-from typing import List
+from typing import List, Dict
 
 from ..general import SysCommand, SysCommandWorker, clear_vt100_escape_codes
 from ..disk.partition import Partition
@@ -16,6 +18,21 @@ class Fido2Device:
 	manufacturer: str
 	product: str
 
+	def json(self) -> Dict[str, str]:
+		return {
+			'path': str(self.path),
+			'manufacturer': self.manufacturer,
+			'product': self.product
+		}
+
+	@classmethod
+	def parse_arg(cls, arg: Dict[str, str]) -> 'Fido2Device':
+		return Fido2Device(
+			Path(arg['path']),
+			arg['manufacturer'],
+			arg['product']
+		)
+
 
 class Fido2:
 	_loaded: bool = False
diff --git a/archinstall/lib/installer.py b/archinstall/lib/installer.py
index 8c6a8367..0e9f0662 100644
--- a/archinstall/lib/installer.py
+++ b/archinstall/lib/installer.py
@@ -248,7 +248,7 @@ class Installer:
 
 		# we manage the encrypted partititons
 		if self._disk_encryption:
-			for partition in self._disk_encryption.partitions:
+			for partition in self._disk_encryption.all_partitions:
 				# open the luks device and all associate stuff
 				loopdev = f"{storage.get('ENC_IDENTIFIER', 'ai')}{pathlib.Path(partition['device_instance'].path).name}"
 
@@ -324,7 +324,7 @@ class Installer:
 			file = f"/{file}"
 		if len(file.strip()) <= 0 or file == '/':
 			raise ValueError(f"The filename for the swap file has to be a valid path, not: {self.target}{file}")
-			
+
 		SysCommand(f'dd if=/dev/zero of={self.target}{file} bs={size} count=1')
 		SysCommand(f'chmod 0600 {self.target}{file}')
 		SysCommand(f'mkswap {self.target}{file}')
@@ -452,7 +452,7 @@ class Installer:
 			if hasattr(plugin, 'on_genfstab'):
 				if plugin.on_genfstab(self) is True:
 					break
-		
+
 		with open(f"{self.target}/etc/fstab", 'a') as fstab_fh:
 			for entry in self.FSTAB_ENTRIES:
 				fstab_fh.write(f'{entry}\n')
diff --git a/archinstall/lib/menu/global_menu.py b/archinstall/lib/menu/global_menu.py
index f0062b4c..7c5b153e 100644
--- a/archinstall/lib/menu/global_menu.py
+++ b/archinstall/lib/menu/global_menu.py
@@ -279,8 +279,8 @@ class GlobalMenu(AbstractMenu):
 			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.all_partitions:
+				output += 'Partitions: {} selected'.format(len(encryption.all_partitions)) + '\n'
 
 			if encryption.hsm_device:
 				output += f'HSM: {encryption.hsm_device.manufacturer}'
diff --git a/archinstall/lib/models/disk_encryption.py b/archinstall/lib/models/disk_encryption.py
index 3edab93e..a4a501d9 100644
--- a/archinstall/lib/models/disk_encryption.py
+++ b/archinstall/lib/models/disk_encryption.py
@@ -1,5 +1,7 @@
+from __future__ import annotations
+
 from dataclasses import dataclass, field
-from enum import Enum, auto
+from enum import Enum
 from typing import Optional, List, Dict, TYPE_CHECKING, Any
 
 from ..hsm.fido import Fido2Device
@@ -9,8 +11,7 @@ if TYPE_CHECKING:
 
 
 class EncryptionType(Enum):
-	Partition = auto()
-	# FullDiskEncryption = auto()
+	Partition = 'partition'
 
 	@classmethod
 	def _encryption_type_mapper(cls) -> Dict[str, 'EncryptionType']:
@@ -35,9 +36,55 @@ class EncryptionType(Enum):
 class DiskEncryption:
 	encryption_type: EncryptionType = EncryptionType.Partition
 	encryption_password: str = ''
-	partitions: List[str] = field(default_factory=list)
+	partitions: Dict[str, List[Dict[str, Any]]] = field(default_factory=dict)
 	hsm_device: Optional[Fido2Device] = None
 
+	@property
+	def all_partitions(self) -> List[Dict[str, Any]]:
+		_all: List[Dict[str, Any]] = []
+		for parts in self.partitions.values():
+			_all += parts
+		return _all
+
 	def generate_encryption_file(self, partition) -> bool:
-		return partition in self.partitions and partition['mountpoint'] != '/'
-	
+		return partition in self.all_partitions and partition['mountpoint'] != '/'
+
+	def json(self) -> Dict[str, Any]:
+		obj = {
+			'encryption_type': self.encryption_type.value,
+			'partitions': self.partitions
+		}
+
+		if self.hsm_device:
+			obj['hsm_device'] = self.hsm_device.json()
+
+		return obj
+
+	@classmethod
+	def parse_arg(
+		cls,
+		disk_layout: Dict[str, Any],
+		arg: Dict[str, Any],
+		password: str = ''
+	) -> 'DiskEncryption':
+		# we have to map the enc partition config to the disk layout objects
+		# they both need to point to the same object as it will get modified
+		# during the installation process
+		enc_partitions: Dict[str, List[Dict[str, Any]]] = {}
+
+		for path, partitions in disk_layout.items():
+			conf_partitions = arg['partitions'].get(path, [])
+			for part in partitions['partitions']:
+				if part in conf_partitions:
+					enc_partitions.setdefault(path, []).append(part)
+
+		enc = DiskEncryption(
+			EncryptionType(arg['encryption_type']),
+			password,
+			enc_partitions
+		)
+
+		if hsm := arg.get('hsm_device', None):
+			enc.hsm_device = Fido2Device.parse_arg(hsm)
+
+		return enc
diff --git a/archinstall/lib/user_interaction/system_conf.py b/archinstall/lib/user_interaction/system_conf.py
index 42a6cec7..e1581677 100644
--- a/archinstall/lib/user_interaction/system_conf.py
+++ b/archinstall/lib/user_interaction/system_conf.py
@@ -60,7 +60,6 @@ def select_harddrives(preset: List[str] = []) -> List[str]:
 	selected_harddrive = Menu(
 		title,
 		list(options.keys()),
-		preset_values=preset,
 		multi=True,
 		allow_reset=True,
 		allow_reset_warning_msg=warning