From 3f0095dd9a90ba51b7cd9087f43161e2e8e0c90e Mon Sep 17 00:00:00 2001
From: Michael Tremer <michael.tremer@ipfire.org>
Date: Sat, 8 May 2021 11:12:47 +0000
Subject: [PATCH] disk: Add automatic partitioning steps

Signed-off-by: Michael Tremer <michael.tremer@ipfire.org>
---
 src/python/__init__.py |   7 +++
 src/python/disk.py     | 109 ++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 115 insertions(+), 1 deletion(-)

diff --git a/src/python/__init__.py b/src/python/__init__.py
index 81a7bb5..724aca3 100644
--- a/src/python/__init__.py
+++ b/src/python/__init__.py
@@ -49,6 +49,9 @@ class Bricklayer(object):
 		# Settings
 		self.settings = {
 			"language" : lang.default_language,
+
+			# Set the default swap size to 1 GiB
+			"swap-size": 1024 ** 3,
 		}
 
 		# Read OS information
@@ -72,6 +75,10 @@ class Bricklayer(object):
 		lang.SelectLanguage,
 		step.Welcome,
 		disk.SelectDisk,
+		disk.CalculatePartitionLayout,
+
+		# Go!
+		disk.CreatePartitionLayout,
 
 		# Done!
 		step.Congratulations,
diff --git a/src/python/disk.py b/src/python/disk.py
index b759805..9869fc2 100644
--- a/src/python/disk.py
+++ b/src/python/disk.py
@@ -29,6 +29,8 @@ from .i18n import _
 # Setup logging
 log = logging.getLogger("bricklayer.disk")
 
+DEFAULT_FILESYSTEM = "ext4"
+
 TEST_DISK_SIZE = 4294967296 # 4GiB
 
 class Disks(object):
@@ -94,6 +96,16 @@ class Disks(object):
 		"""
 		return [disk for disk in self.disks if disk.selected]
 
+	def calculate_partition_layout(self):
+		"""
+			This creates the partition layout, but doesn't write it to disk, yet
+		"""
+		# Find the root device
+		root = self.selected[0] # XXX select the first harddisk only
+
+		# Create one giant root partition
+		root.create_system_partitions()
+
 
 class Disk(object):
 	def __init__(self, bricklayer, device):
@@ -102,11 +114,17 @@ class Disk(object):
 		# The parted device
 		self.device = device
 
+		# The parted disk (with a blank partition table)
+		self.parted = parted.freshDisk(self.device, "gpt")
+
 		# Has this device been selected?
 		self.selected = False
 
+		# Where are we starting with the next partition?
+		self._start = 1
+
 	def __repr__(self):
-		return "<%s %s>" % (self.__class__.__name__, self)
+		return "<%s %s>" % (self.__class__.__name__, self.path)
 
 	def __str__(self):
 		return "%s (%s)" % (self.model, util.format_size(self.size))
@@ -153,6 +171,73 @@ class Disk(object):
 	def size(self):
 		return self.device.length * self.device.sectorSize
 
+	def create_system_partitions(self):
+		"""
+			This method creates a basic partition layout on this disk with all
+			partitions that the systems needs. This is as follows:
+
+			1) BIOS boot partition (used for GRUB, etc.)
+			2) A FAT-formatted EFI partition
+			3) A swap partition
+			4) A / partition
+		"""
+		log.debug("Creating partition layout on %s" % self.path)
+
+		# Create a bootloader partition of exactly 1 MiB
+		self._add_partition("BOOTLDR", DEFAULT_FILESYSTEM, length=1024**2,
+			flags=[parted.PARTITION_BIOS_GRUB])
+
+		# Create an EFI-partition of exactly 32 MiB
+		self._add_partition("ESP", "fat32", length=32 * 1024**2,
+			flags=[parted.PARTITION_ESP])
+
+		# Create a swap partition
+		swap_size = self.bricklayer.settings.get("swap-size", 0)
+		if swap_size:
+			self._add_partition("SWAP", "linux-swap(v1)", length=swap_size)
+
+		# Use all remaining space for root
+		self._add_partition("ROOT", DEFAULT_FILESYSTEM)
+
+	def _add_partition(self, name, filesystem, length=None, flags=[]):
+		# The entire device
+		if length is None:
+			length = self.device.getLength() - self._start
+
+		# Otherwise convert bytes into sectors
+		else:
+			length = length // self.device.sectorSize
+
+		# Calculate the partitions geometry
+		geometry = parted.Geometry(self.device, start=self._start, length=length)
+
+		# Create the filesystem
+		fs = parted.FileSystem(type=filesystem, geometry=geometry)
+
+		# Create the partition
+		partition = parted.Partition(disk=self.parted, type=parted.PARTITION_NORMAL,
+			fs=fs, geometry=geometry)
+
+		# Set name
+		partition.name = name
+
+		# Set flags
+		for flag in flags:
+			partition.setFlag(flag)
+
+		# Add the partition and align the most optimal way
+		self.parted.addPartition(partition,
+			constraint=self.device.optimalAlignedConstraint)
+
+		# Store end to know where to begin the next partition
+		self._start = geometry.end + 1
+
+	def commit(self):
+		"""
+			Write the partition table to disk
+		"""
+		self.parted.commit()
+
 
 class SelectDisk(step.InteractiveStep):
 	"""
@@ -202,3 +287,25 @@ class SelectDisk(step.InteractiveStep):
 				disk.selected = disk in selection
 
 			break
+
+
+class CalculatePartitionLayout(step.Step):
+	"""
+		Calculates the partition layout
+	"""
+	def run(self, tui):
+		# This probably will be fast enough that we do not need to show anything
+
+		# Perform the job
+		self.bricklayer.disks.calculate_partition_layout()
+
+
+class CreatePartitionLayout(step.Step):
+	"""
+		Creates the desired partition layout on disk
+	"""
+	def run(self, tui):
+		log.debug("Creating partitions")
+
+		for disk in self.bricklayer.disks.selected:
+			disk.commit()
-- 
2.47.2