2 files changed, 45 insertions, 16 deletions

diff --git a/archinstall/lib/disk.py b/archinstall/lib/disk.py
index 01a736a3..3397d7cb 100644
--- a/archinstall/lib/disk.py
+++ b/archinstall/lib/disk.py
@@ -197,6 +197,21 @@ class Partition():
 
 		return True
 
+	def encrypt(self, *args, **kwargs):
+		"""
+		A wrapper function for luks2() instances and the .encrypt() method of that instance.
+		"""
+		from .luks import luks2
+
+		if not self.encrypted:
+			raise DiskError(f"Attempting to encrypt a partition that was not marked for encryption: {self}")
+
+		if not self.safe_to_format():
+			return False
+
+		handle = luks2(self, None, None)
+		return handle.encrypt(self, *args, **kwargs)
+
 	def format(self, filesystem=None, path=None, allow_formatting=None, log_formating=True):
 		"""
 		Format can be given an overriding path, for instance /dev/null to test
@@ -204,6 +219,7 @@ class Partition():
 		"""
 		if filesystem is None:
 			filesystem = self.filesystem
+
 		if path is None:
 			path = self.path
 		if allow_formatting is None:
@@ -349,15 +365,23 @@ class Filesystem():
 		"""
 		return self.raw_parted(string).exit_code
 
-	def use_entire_disk(self, prep_mode=None):
+	def use_entire_disk(self, root_filesystem_type='ext4', encrypt_root_partition=True):
 		self.add_partition('primary', start='1MiB', end='513MiB', format='vfat')
 		self.set_name(0, 'EFI')
 		self.set(0, 'boot on')
-		self.set(0, 'esp on') # TODO: Redundant, as in GPT mode it's an alias for "boot on"? https://www.gnu.org/software/parted/manual/html_node/set.html
-		if prep_mode == 'luks2':
-			self.add_partition('primary', start='513MiB', end='100%')
-		else:
-			self.add_partition('primary', start='513MiB', end='100%', format=prep_mode)
+		# TODO: Probably redundant because in GPT mode 'esp on' is an alias for "boot on"?
+		# https://www.gnu.org/software/parted/manual/html_node/set.html
+		self.set(0, 'esp on')
+		self.add_partition('primary', start='513MiB', end='100%')
+
+		self.blockdevice.partition[0].filesystem = 'vfat'
+		self.blockdevice.partition[1].filesystem = root_filesystem_type
+
+		self.blockdevice.partition[0].target_mountpoint = '/boot'
+		self.blockdevice.partition[1].target_mountpoint = '/'
+
+		if encrypt_root_partition:
+			self.blockdevice.partition[1].encrypted = True
 
 	def add_partition(self, type, start, end, format=None):
 		log(f'Adding partition to {self.blockdevice}', level=LOG_LEVELS.Info)
diff --git a/archinstall/lib/luks.py b/archinstall/lib/luks.py
index b98994ef..c9946239 100644
--- a/archinstall/lib/luks.py
+++ b/archinstall/lib/luks.py
@@ -16,16 +16,17 @@ class luks2():
 		self.filesystem = 'crypto_LUKS'
 
 	def __enter__(self):
-		if self.partition.allow_formatting:
-			self.key_file = self.encrypt(self.partition, *self.args, **self.kwargs)
-		else:
-			if not self.key_file:
-				self.key_file = f"/tmp/{os.path.basename(self.partition.path)}.disk_pw"  # TODO: Make disk-pw-file randomly unique?
-			
-			if type(self.password) != bytes: self.password = bytes(self.password, 'UTF-8')
+		#if self.partition.allow_formatting:
+		#	self.key_file = self.encrypt(self.partition, *self.args, **self.kwargs)
+		#else:
+		if not self.key_file:
+			self.key_file = f"/tmp/{os.path.basename(self.partition.path)}.disk_pw"  # TODO: Make disk-pw-file randomly unique?
+		
+		if type(self.password) != bytes:
+			self.password = bytes(self.password, 'UTF-8')
 
-			with open(self.key_file, 'wb') as fh:
-				fh.write(self.password)
+		with open(self.key_file, 'wb') as fh:
+			fh.write(self.password)
 
 		return self.unlock(self.partition, self.mountpoint, self.key_file)
 
@@ -38,6 +39,9 @@ class luks2():
 	def encrypt(self, partition, password=None, key_size=512, hash_type='sha512', iter_time=10000, key_file=None):
 		# TODO: We should be able to integrate this into the main log some how.
 		#       Perhaps post-mortem?
+		if not self.partition.allow_formatting:
+			raise DiskError(f'Could not encrypt volume {self.partition} due to it having a formatting lock.')
+
 		log(f'Encrypting {partition} (This might take a while)', level=LOG_LEVELS.Info)
 
 		if not key_file:
@@ -49,7 +53,8 @@ class luks2():
 		if not password:
 			password = self.password
 
-		if type(password) != bytes: password = bytes(password, 'UTF-8')
+		if type(password) != bytes:
+			password = bytes(password, 'UTF-8')
 
 		with open(key_file, 'wb') as fh:
 			fh.write(password)