Support encrypting multiple partitions (#759)

* Added support for storing disk encryption keyfiles and add them to a keyslot.
* Added a luks2().add_key() function in order to inject a keyfile into a keyslot on a encrypted volume.
* Simplified 'missing encryption password' logic in Filesystem(). Added a call to luks2().add_key() after the root-password is set on the volume, to add the keyfile in slot 2
* Adding in password handling in luks2().add_key(). It's required to enter a previous passphrase to unlock the volume and add a new keyslot. Also simplified the handling of partition in Installer().mount_ordered_layout()
* Adding in encryption on all partitions except /boot when encryption is opted in
* Removed setting size on Partition() as it's a read only value. No idea how Partition().size = size hasn't caused an issue before. Removed size=X argument to Partition()
* Added a uniqueness to the loopdevice name. This should ensure that multiple encrypted volumes can be opened at the same time, except for Partition().detect_inner_filesystem() operations which can only happen one at a time since they share namespace. This should never be an issue since archinstall is single threaded and no concurrent operations can/should happen.
* Added partprobe() as part of disk/helpers.py, added a /dev/ -> UUID mapper function called convert_device_to_uuid(path). Added a luks2().crypttab() function that sets up a /etc/crypttab entry.
* Moved the responsability for telling archinstall to generate a keyfile from Filesystem() to user_interaction.py. This should in the future be a user-input based value, and not something the Filesystem() automatically dictates.
* Added a retry mechanism to luks2().encrypt() to avoid having to re-start the installation when a device simply wasn't up yet.
* Swapping UUID= lookup from loopdev to physdev.

8 files changed, 114 insertions, 41 deletions

diff --git a/archinstall/lib/disk/blockdevice.py b/archinstall/lib/disk/blockdevice.py
index d8c34893..f8575de4 100644
--- a/archinstall/lib/disk/blockdevice.py
+++ b/archinstall/lib/disk/blockdevice.py
@@ -128,7 +128,7 @@ class BlockDevice:
 				if part_id not in self.part_cache:
 					# TODO: Force over-write even if in cache?
 					if part_id not in self.part_cache or self.part_cache[part_id].size != part['size']:
-						self.part_cache[part_id] = Partition(root_path + part_id, self, part_id=part_id, size=part['size'])
+						self.part_cache[part_id] = Partition(root_path + part_id, self, part_id=part_id)
 
 		return {k: self.part_cache[k] for k in sorted(self.part_cache)}
 
@@ -156,7 +156,7 @@ class BlockDevice:
 	@property
 	def size(self):
 		from .helpers import convert_size_to_gb
-		
+
 		output = json.loads(SysCommand(f"lsblk --json -b -o+SIZE {self.path}").decode('UTF-8'))
 
 		for device in output['blockdevices']:
diff --git a/archinstall/lib/disk/filesystem.py b/archinstall/lib/disk/filesystem.py
index 83d7e34f..edf54eb5 100644
--- a/archinstall/lib/disk/filesystem.py
+++ b/archinstall/lib/disk/filesystem.py
@@ -1,6 +1,7 @@
 import time
 import logging
 import json
+import pathlib
 from .partition import Partition
 from .validators import valid_fs_type
 from ..exceptions import DiskError
@@ -80,17 +81,22 @@ class Filesystem:
 
 			if partition.get('filesystem', {}).get('format', False):
 				if partition.get('encrypted', False):
-					if not partition.get('!password') and not storage['arguments'].get('!encryption-password'):
-						if storage['arguments'] == 'silent':
-							raise ValueError(f"Missing encryption password for {partition['device_instance']}")
-						else:
+					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
-							partition['!password'] = get_password(f"Enter a encryption password for {partition['device_instance']}")
-					elif not partition.get('!password') and storage['arguments'].get('!encryption-password'):
+							storage['arguments']['!encryption-password'] = get_password(f"Enter a encryption password for {partition['device_instance']}")
+
 						partition['!password'] = storage['arguments']['!encryption-password']
 
+					loopdev = f"{storage.get('ENC_IDENTIFIER', 'ai')}{pathlib.Path(partition['mountpoint']).name}loop"
+
 					partition['device_instance'].encrypt(password=partition['!password'])
-					with luks2(partition['device_instance'], storage.get('ENC_IDENTIFIER', 'ai') + 'loop', partition['!password']) as unlocked_device:
+
+					# Immediately unlock the encrypted device to format the inner volume
+					with luks2(partition['device_instance'], loopdev, partition['!password'], auto_unmount=True) as unlocked_device:
 						if not partition.get('format'):
 							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 46d86bd5..9442f1b6 100644
--- a/archinstall/lib/disk/helpers.py
+++ b/archinstall/lib/disk/helpers.py
@@ -1,13 +1,15 @@
-import re
-import os
 import json
 import logging
+import os
 import pathlib
+import re
+import time
 from typing import Union
 from .blockdevice import BlockDevice
 from ..exceptions import SysCallError, DiskError
 from ..general import SysCommand
 from ..output import log
+from ..storage import storage
 
 ROOT_DIR_PATTERN = re.compile('^.*?/devices')
 GIGA = 2 ** 30
@@ -209,3 +211,19 @@ def find_partition_by_mountpoint(block_devices, relative_mountpoint :str):
 		for partition in block_devices[device]['partitions']:
 			if partition.get('mountpoint', None) == relative_mountpoint:
 				return partition
+
+def partprobe():
+	SysCommand(f'bash -c "partprobe"')
+
+def convert_device_to_uuid(path :str) -> str:
+	for i in range(storage['DISK_RETRY_ATTEMPTS']):
+		partprobe()
+		output = json.loads(SysCommand(f"lsblk --json -o+UUID {path}").decode('UTF-8'))
+
+		for device in output['blockdevices']:
+			if (dev_uuid := device.get('uuid', None)):
+				return dev_uuid
+
+		time.sleep(storage['DISK_TIMEOUTS'])
+
+	raise DiskError(f"Could not retrieve the UUID of {path} within a timely manner.")
\ No newline at end of file
diff --git a/archinstall/lib/disk/partition.py b/archinstall/lib/disk/partition.py
index b27c8459..d3efe5cf 100644
--- a/archinstall/lib/disk/partition.py
+++ b/archinstall/lib/disk/partition.py
@@ -15,7 +15,7 @@ from ..general import SysCommand
 
 
 class Partition:
-	def __init__(self, path: str, block_device: BlockDevice, part_id=None, size=-1, filesystem=None, mountpoint=None, encrypted=False, autodetect_filesystem=True):
+	def __init__(self, path: str, block_device: BlockDevice, part_id=None, filesystem=None, mountpoint=None, encrypted=False, autodetect_filesystem=True):
 		if not part_id:
 			part_id = os.path.basename(path)
 
@@ -25,7 +25,6 @@ class Partition:
 		self.mountpoint = mountpoint
 		self.target_mountpoint = mountpoint
 		self.filesystem = filesystem
-		self.size = size  # TODO: Refresh?
 		self._encrypted = None
 		self.encrypted = encrypted
 		self.allow_formatting = False
diff --git a/archinstall/lib/general.py b/archinstall/lib/general.py
index a057af5a..48de4cbe 100644
--- a/archinstall/lib/general.py
+++ b/archinstall/lib/general.py
@@ -2,8 +2,10 @@ import hashlib
 import json
 import logging
 import os
+import secrets
 import shlex
 import subprocess
+import string
 import sys
 import time
 from datetime import datetime, date
@@ -46,6 +48,9 @@ from .storage import storage
 def gen_uid(entropy_length=256):
 	return hashlib.sha512(os.urandom(entropy_length)).hexdigest()
 
+def generate_password(length=64):
+	haystack = string.printable # digits, ascii_letters, punctiation (!"#$[] etc) and whitespace
+	return ''.join(secrets.choice(haystack) for i in range(length))
 
 def multisplit(s, splitters):
 	s = [s, ]
@@ -61,7 +66,6 @@ def multisplit(s, splitters):
 		s = ns
 	return s
 
-
 def locate_binary(name):
 	for PATH in os.environ['PATH'].split(':'):
 		for root, folders, files in os.walk(PATH):
diff --git a/archinstall/lib/installer.py b/archinstall/lib/installer.py
index 07de94e0..2629c0f7 100644
--- a/archinstall/lib/installer.py
+++ b/archinstall/lib/installer.py
@@ -8,7 +8,7 @@ import pathlib
 import subprocess
 import glob
 from .disk import get_partitions_in_use, Partition
-from .general import SysCommand
+from .general import SysCommand, generate_password
 from .hardware import has_uefi, is_vm, cpu_vendor
 from .locale_helpers import verify_keyboard_layout, verify_x11_keyboard_layout
 from .disk.helpers import get_mount_info
@@ -187,18 +187,35 @@ class Installer:
 				mountpoints[partition['mountpoint']] = partition
 
 		for mountpoint in sorted(mountpoints.keys()):
-			if mountpoints[mountpoint].get('encrypted', False):
-				loopdev = storage.get('ENC_IDENTIFIER', 'ai') + 'loop'
-				if not (password := mountpoints[mountpoint].get('!password', None)):
-					raise RequirementError(f"Missing mountpoint {mountpoint} encryption password in layout: {mountpoints[mountpoint]}")
+			partition = mountpoints[mountpoint]
+
+			if partition.get('encrypted', False):
+				loopdev = f"{storage.get('ENC_IDENTIFIER', 'ai')}{pathlib.Path(partition['mountpoint']).name}loop"
+				if not (password := partition.get('!password', None)):
+					raise RequirementError(f"Missing mountpoint {mountpoint} encryption password in layout: {partition}")
+
+				with (luks_handle := luks2(partition['device_instance'], loopdev, password, auto_unmount=False)) as unlocked_device:
+					if partition.get('generate-encryption-key-file'):
+						if not (cryptkey_dir := pathlib.Path(f"{self.target}/etc/cryptsetup-keys.d")).exists():
+							cryptkey_dir.mkdir(parents=True, exist_ok=True)
+
+						# Once we store the key as ../xyzloop.key systemd-cryptsetup can automatically load this key
+						# if we name the device to "xyzloop".
+						encryption_key_path = f"/etc/cryptsetup-keys.d/{pathlib.Path(partition['mountpoint']).name}loop.key"
+						with open(f"{self.target}{encryption_key_path}", "w") as keyfile:
+							keyfile.write(generate_password(length=512))
+
+						os.chmod(encryption_key_path, 0o400)
+
+						luks_handle.add_key(pathlib.Path(f"{self.target}{encryption_key_path}"), password=password)
+						luks_handle.crypttab(self, encryption_key_path, options=["luks", "key-slot=1"])
 
-				with luks2(mountpoints[mountpoint]['device_instance'], loopdev, password, auto_unmount=False) as unlocked_device:
 					log(f"Mounting {mountpoint} to {self.target}{mountpoint} using {unlocked_device}", level=logging.INFO)
 					unlocked_device.mount(f"{self.target}{mountpoint}")
 
 			else:
-				log(f"Mounting {mountpoint} to {self.target}{mountpoint} using {mountpoints[mountpoint]['device_instance']}", level=logging.INFO)
-				mountpoints[mountpoint]['device_instance'].mount(f"{self.target}{mountpoint}")
+				log(f"Mounting {mountpoint} to {self.target}{mountpoint} using {partition['device_instance']}", level=logging.INFO)
+				partition['device_instance'].mount(f"{self.target}{mountpoint}")
 
 			time.sleep(1)
 			try:
@@ -206,7 +223,7 @@ class Installer:
 			except DiskError:
 				raise DiskError(f"Target {self.target}{mountpoint} never got mounted properly (unable to get mount information using findmnt).")
 
-			if (subvolumes := mountpoints[mountpoint].get('btrfs', {}).get('subvolumes', {})):
+			if (subvolumes := partition.get('btrfs', {}).get('subvolumes', {})):
 				for name, location in subvolumes.items():
 					create_subvolume(self, location)
 					mount_subvolume(self, location)
diff --git a/archinstall/lib/luks.py b/archinstall/lib/luks.py
index b93403ef..55eaa62f 100644
--- a/archinstall/lib/luks.py
+++ b/archinstall/lib/luks.py
@@ -4,11 +4,11 @@ import os
 import pathlib
 import shlex
 import time
-
-from .disk import Partition
-from .general import SysCommand
+from .disk import Partition, convert_device_to_uuid
+from .general import SysCommand, SysCommandWorker
 from .output import log
 from .exceptions import SysCallError, DiskError
+from .storage import storage
 
 class luks2:
 	def __init__(self, partition, mountpoint, password, key_file=None, auto_unmount=False, *args, **kwargs):
@@ -78,8 +78,17 @@ class luks2:
 		])
 
 		try:
-			# Try to setup the crypt-device
-			cmd_handle = SysCommand(cryptsetup_args)
+			# Retry formatting the volume because archinstall can some times be too quick
+			# which generates a "Device /dev/sdX does not exist or access denied." between
+			# setting up partitions and us trying to encrypt it.
+			for i in range(storage['DISK_RETRY_ATTEMPTS']):
+				if (cmd_handle := SysCommand(cryptsetup_args)).exit_code != 0:
+					time.sleep(storage['DISK_TIMEOUTS'])
+				else:
+					break
+
+			if cmd_handle.exit_code != 0:
+				raise DiskError(f'Could not encrypt volume "{partition.path}": {b"".join(cmd_handle)}')
 		except SysCallError as err:
 			if err.exit_code == 256:
 				log(f'{partition} is being used, trying to unmount and crypt-close the device and running one more attempt at encrypting the device.', level=logging.DEBUG)
@@ -108,9 +117,6 @@ class luks2:
 			else:
 				raise err
 
-		if cmd_handle.exit_code != 0:
-			raise DiskError(f'Could not encrypt volume "{partition.path}": {b"".join(cmd_handle)}')
-
 		return key_file
 
 	def unlock(self, partition, mountpoint, key_file):
@@ -146,3 +152,24 @@ class luks2:
 	def format(self, path):
 		if (handle := SysCommand(f"/usr/bin/cryptsetup -q -v luksErase {path}")).exit_code != 0:
 			raise DiskError(f'Could not format {path} with {self.filesystem} because: {b"".join(handle)}')
+
+	def add_key(self, path :pathlib.Path, password :str):
+		if not path.exists():
+			raise OSError(2, f"Could not import {path} as a disk encryption key, file is missing.", str(path))
+
+		log(f'Adding additional key-file {path} for {self.partition}', level=logging.INFO)
+
+		worker = SysCommandWorker(f"/usr/bin/cryptsetup -q -v luksAddKey {self.partition.path} {path}")
+		pw_injected = False
+		while worker.is_alive():
+			if b'Enter any existing passphrase' in worker and pw_injected is False:
+				worker.write(bytes(password, 'UTF-8'))
+				pw_injected = True
+
+		if worker.exit_code != 0:
+			raise DiskError(f'Could not add encryption key {path} to {self.partition} because: {worker}')
+
+	def crypttab(self, installation, key_path :str, options=["luks", "key-slot=1"]):
+		log(f'Adding a crypttab entry for key {key_path} in {installation}', level=logging.INFO)
+		with open(f"{installation.target}/etc/crypttab", "a") as crypttab:
+			crypttab.write(f"{self.mountpoint} UUID={convert_device_to_uuid(self.partition.path)} {key_path} {','.join(options)}\n")
\ No newline at end of file
diff --git a/archinstall/lib/user_interaction.py b/archinstall/lib/user_interaction.py
index 39e87c02..7d90915f 100644
--- a/archinstall/lib/user_interaction.py
+++ b/archinstall/lib/user_interaction.py
@@ -9,7 +9,7 @@ import signal
 import sys
 import time
 
-from .disk import BlockDevice, valid_fs_type, find_partition_by_mountpoint, suggest_single_disk_layout, suggest_multi_disk_layout, valid_parted_position
+from .disk import BlockDevice, valid_fs_type, suggest_single_disk_layout, suggest_multi_disk_layout, valid_parted_position
 from .exceptions import RequirementError, UserError, DiskError
 from .hardware import AVAILABLE_GFX_DRIVERS, has_uefi, has_amd_graphics, has_intel_graphics, has_nvidia_graphics
 from .locale_helpers import list_keyboard_languages, verify_keyboard_layout, search_keyboard_layout
@@ -193,18 +193,20 @@ def generic_multi_select(options, text="Select one or more of the options above
 	return selected_options
 
 def select_encrypted_partitions(block_devices :dict, password :str) -> dict:
-	root = find_partition_by_mountpoint(block_devices, '/')
-	root['encrypted'] = True
-	root['!password'] = password
+	for device in block_devices:
+		for partition in block_devices[device]['partitions']:
+			if partition.get('mountpoint', None) != '/boot':
+				partition['encrypted'] = True
+				partition['!password'] = password
 
-	return block_devices
+				if partition['mountpoint'] != '/':
+					# Tell the upcoming steps to generate a key-file for non root mounts.
+					partition['generate-encryption-key-file'] = True
 
-	# TODO: Next version perhaps we can support multiple encrypted partitions
-	# options = []
-	# for partition in block_devices.values():
-	# 	options.append({key: val for key, val in partition.items() if val})
+	return block_devices
 
-	# print(generic_multi_select(options, f"Choose which partitions to encrypt (leave blank when done): "))
+	# TODO: Next version perhaps we can support mixed multiple encrypted partitions
+	# Users might want to single out a partition for non-encryption to share between dualboot etc.
 
 class MiniCurses:
 	def __init__(self, width, height):