1 files changed, 114 insertions, 4 deletions

diff --git a/archinstall/lib/disk.py b/archinstall/lib/disk.py
index 8f67111a..ac59600d 100644
--- a/archinstall/lib/disk.py
+++ b/archinstall/lib/disk.py
@@ -31,10 +31,10 @@ class BlockDevice:
 		self.info = info
 		self.keep_partitions = True
 		self.part_cache = OrderedDict()
+
 		# TODO: Currently disk encryption is a BIT misleading.
 		#       It's actually partition-encryption, but for future-proofing this
 		#       I'm placing the encryption password on a BlockDevice level.
-		self.encryption_password = None
 
 	def __repr__(self, *args, **kwargs):
 		return f"BlockDevice({self.device})"
@@ -48,6 +48,9 @@ class BlockDevice:
 			raise KeyError(f'{self} does not contain information: "{key}"')
 		return self.info[key]
 
+	def __len__(self):
+		return len(self.partitions)
+
 	def json(self):
 		"""
 		json() has precedence over __dump__, so this is a way
@@ -61,12 +64,22 @@ class BlockDevice:
 
 	def __dump__(self):
 		return {
-			'path': self.path,
-			'info': self.info,
-			'partition_cache': self.part_cache
+			self.path : {
+				'partuuid' : self.uuid,
+				'wipe' : self.info.get('wipe', None),
+				'partitions' : [part.__dump__() for part in self.partitions.values()]
+			}
 		}
 
 	@property
+	def partition_type(self):
+		output = b"".join(sys_command(f"lsblk --json -o+PTTYPE {self.path}"))
+		output = json.loads(output.decode('UTF-8'))
+	
+		for device in output['blockdevices']:
+			return device['pttype']
+
+	@property
 	def device(self):
 		"""
 		Returns the actual device-endpoint of the BlockDevice.
@@ -144,6 +157,16 @@ class BlockDevice:
 		for partition in json.loads(lsblk.decode('UTF-8'))['blockdevices']:
 			return partition.get('uuid', None)
 
+	@property
+	def size(self):
+		output = b"".join(sys_command(f"lsblk --json -o+SIZE {self.path}"))
+		output = json.loads(output.decode('UTF-8'))
+	
+		for device in output['blockdevices']:
+			assert device['size'][-1] == 'G' # Make sure we're counting in Gigabytes, otherwise the next logic fails.
+
+			return float(device['size'][:-1])
+
 	def has_partitions(self):
 		return len(self.partitions)
 
@@ -210,6 +233,82 @@ class Partition:
 		else:
 			return f'Partition(path={self.path}, size={self.size}, fs={self.filesystem}{mount_repr})'
 
+	def __dump__(self):
+		return {
+			'type' : 'primary',
+			'PARTUUID' : self.uuid,
+			'wipe' : self.allow_formatting,
+			'boot' : self.boot,
+			'ESP' : self.boot,
+			'mountpoint' : self.target_mountpoint,
+			'encrypted' : self._encrypted,
+			'start' : self.start,
+			'size' : self.end,
+			'filesystem' : {
+				'format' : get_filesystem_type(self.path)
+			}
+		}
+
+	@property
+	def sector_size(self):
+		output = b"".join(sys_command(f"lsblk --json -o+LOG-SEC {self.path}"))
+		output = json.loads(output.decode('UTF-8'))
+		
+		for device in output['blockdevices']:
+			return device.get('log-sec', None)
+
+	@property
+	def start(self):
+		output = b"".join(sys_command(f"sfdisk --json {self.block_device.path}"))
+		output = json.loads(output.decode('UTF-8'))
+	
+		for partition in output.get('partitiontable', {}).get('partitions', []):
+			if partition['node'] == self.path:
+				return partition['start']# * self.sector_size
+
+	@property
+	def end(self):
+		# TODO: Verify that the logic holds up, that 'size' is the size without 'start' added to it.
+		output = b"".join(sys_command(f"sfdisk --json {self.block_device.path}"))
+		output = json.loads(output.decode('UTF-8'))
+
+		for partition in output.get('partitiontable', {}).get('partitions', []):
+			if partition['node'] == self.path:
+				return partition['size']# * self.sector_size
+
+	@property
+	def boot(self):
+		output = b"".join(sys_command(f"sfdisk --json {self.block_device.path}"))
+		output = json.loads(output.decode('UTF-8'))
+
+		# Get the bootable flag from the sfdisk output:
+		# {
+		#    "partitiontable": {
+		#       "label":"dos",
+		#       "id":"0xd202c10a",
+		#       "device":"/dev/loop0",
+		#       "unit":"sectors",
+		#       "sectorsize":512,
+		#       "partitions": [
+		#          {"node":"/dev/loop0p1", "start":2048, "size":10483712, "type":"83", "bootable":true}
+		#       ]
+		#    }
+		# }
+
+		for partition in output.get('partitiontable', {}).get('partitions', []):
+			if partition['node'] == self.path:
+				return partition.get('bootable', False)
+
+		return False
+
+	@property
+	def partition_type(self):
+		output = b"".join(sys_command(f"lsblk --json -o+PTTYPE {self.path}"))
+		output = json.loads(output.decode('UTF-8'))
+	
+		for device in output['blockdevices']:
+			return device['pttype']
+
 	@property
 	def uuid(self) -> Optional[str]:
 		"""
@@ -663,3 +762,14 @@ def disk_layouts():
 	except SysCallError as err:
 		log(f"Could not return disk layouts: {err}")
 		return None
+
+
+def encrypted_partitions(blockdevices :dict) -> bool:
+	for partition in blockdevices.values():
+		if partition.get('encrypted', False):
+			yield partition
+
+def find_partition_by_mountpoint(partitions, relative_mountpoint :str):
+	for partition in partitions:
+		if partition.get('mountpoint', None) == relative_mountpoint:
+			return partition
\ No newline at end of file