1 files changed, 53 insertions, 6 deletions

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