From f19f35897b6bcbf5d510d6e7cd672332f86a1f3a Mon Sep 17 00:00:00 2001
From: codefiles <11915375+codefiles@users.noreply.github.com>
Date: Tue, 19 Nov 2024 02:37:24 -0500
Subject: [PATCH] Add _get_key_file() (#2882)

---
 archinstall/lib/luks.py | 36 +++++++++++++-----------------------
 1 file changed, 13 insertions(+), 23 deletions(-)

diff --git a/archinstall/lib/luks.py b/archinstall/lib/luks.py
index 8d2b558b..ad2ab91a 100644
--- a/archinstall/lib/luks.py
+++ b/archinstall/lib/luks.py
@@ -53,9 +53,6 @@ class Luks2:
 		if self.auto_unmount:
 			self.lock()
 
-	def _default_key_file(self) -> Path:
-		return Path(f'/tmp/{self.luks_dev_path.name}.disk_pw')
-
 	def _password_bytes(self) -> bytes:
 		if not self.password:
 			raise ValueError('Password for luks2 device was not specified')
@@ -65,6 +62,17 @@ class Luks2:
 		else:
 			return bytes(self.password, 'UTF-8')
 
+	def _get_key_file(self, key_file: Path | None = None) -> Path:
+		if key_file:
+			return key_file
+
+		if self.key_file:
+			return self.key_file
+
+		default_key_file = Path(f'/tmp/{self.luks_dev_path.name}.disk_pw')
+		default_key_file.write_bytes(self._password_bytes())
+		return default_key_file
+
 	def encrypt(
 		self,
 		key_size: int = 512,
@@ -74,16 +82,7 @@ class Luks2:
 	) -> Path:
 		debug(f'Luks2 encrypting: {self.luks_dev_path}')
 
-		byte_password = self._password_bytes()
-
-		if not key_file:
-			if self.key_file:
-				key_file = self.key_file
-			else:
-				key_file = self._default_key_file()
-
-				with open(key_file, 'wb') as fh:
-					fh.write(byte_password)
+		key_file = self._get_key_file(key_file)
 
 		cryptsetup_args = shlex.join([
 			'/usr/bin/cryptsetup',
@@ -155,16 +154,7 @@ class Luks2:
 		if not self.mapper_name:
 			raise ValueError('mapper name missing')
 
-		byte_password = self._password_bytes()
-
-		if not key_file:
-			if self.key_file:
-				key_file = self.key_file
-			else:
-				key_file = self._default_key_file()
-
-				with open(key_file, 'wb') as fh:
-					fh.write(byte_password)
+		key_file = self._get_key_file(key_file)
 
 		wait_timer = time.time()
 		while Path(self.luks_dev_path).exists() is False and time.time() - wait_timer < 10: