Attempting to select drives by size and/or id

3 files changed, 31 insertions, 2 deletions

diff --git a/archinstall/lib/disk.py b/archinstall/lib/disk.py
index e50f1b08..2f3c122f 100644
--- a/archinstall/lib/disk.py
+++ b/archinstall/lib/disk.py
@@ -34,7 +34,8 @@ class BlockDevice():
 		"""
 		return {
 			'path' : self.path,
-			'size' : self.info['size'] if 'size' in self.info else '<unknown>'
+			'size' : self.info['size'] if 'size' in self.info else '<unknown>',
+			'model' : self.info['model'] if 'model' in self.info else '<unknown>'
 		}
 
 	def __dump__(self):
@@ -252,8 +253,29 @@ def all_disks(*args, **kwargs):
 	kwargs.setdefault("partitions", False)
 	drives = OrderedDict()
 	#for drive in json.loads(sys_command(f'losetup --json', *args, **lkwargs, hide_from_log=True)).decode('UTF_8')['loopdevices']:
-	for drive in json.loads(b''.join(sys_command(f'lsblk --json -l -n -o path,size,type,mountpoint,label,pkname', *args, **kwargs, hide_from_log=True)).decode('UTF_8'))['blockdevices']:
+	for drive in json.loads(b''.join(sys_command(f'lsblk --json -l -n -o path,size,type,mountpoint,label,pkname,model', *args, **kwargs, hide_from_log=True)).decode('UTF_8'))['blockdevices']:
 		if not kwargs['partitions'] and drive['type'] == 'part': continue
 
 		drives[drive['path']] = BlockDevice(drive['path'], drive)
 	return drives
+
+def convert_to_gigabytes(string):
+	unit = string.strip()[-1]
+	size = float(string.strip()[:-1])
+
+	if unit == 'M':
+		size = size/1024
+	elif unit == 'T':
+		size = size*1024
+
+	return size
+
+def harddrive(size=None, model=None, fuzzy=False):
+	collection = all_disks()
+	for drive in collection:
+		if size and convert_to_gigabytes(collection[drive]['size']) != size:
+			continue
+		if model and (collection[drive]['model'] is None or collection[drive]['model'].lower() != model.lower()):
+			continue
+
+		return drive
\ No newline at end of file
diff --git a/archinstall/lib/installer.py b/archinstall/lib/installer.py
index 543f2ca3..86eafc88 100644
--- a/archinstall/lib/installer.py
+++ b/archinstall/lib/installer.py
@@ -89,6 +89,10 @@ class Installer():
 			self.log(f"Submit this zip file as an issue to https://github.com/Torxed/archinstall/issues", level=LOG_LEVELS.Warning)
 			return False
 
+	def mount(self, mountpoint, partition):
+		partition.mount(f'{self.mountpoint}/srv/http')
+
+
 	def post_install_check(self, *args, **kwargs):
 		return [step for step, flag in self.helper_flags.items() if flag is False]
 
diff --git a/test.py b/test.py
new file mode 100644
index 00000000..61ba1925
--- /dev/null
+++ b/test.py
@@ -0,0 +1,3 @@
+import archinstall
+
+print(archinstall.harddrive(size=111.8, model='KINGSTON_SKC100S3120G'))