disks: Fix automatic disk selection in unattended mode

[people/ms/bricklayer.git] / src / python / disk.py

###############################################################################
#                                                                             #
# Bricklayer - An Installer for IPFire                                        #
# Copyright (C) 2021 IPFire Development Team                                  #
#                                                                             #
# This program is free software; you can redistribute it and/or               #
# modify it under the terms of the GNU General Public License                 #
# as published by the Free Software Foundation; either version 2              #
# of the License, or (at your option) any later version.                      #
#                                                                             #
# This program is distributed in the hope that it will be useful,             #
# but WITHOUT ANY WARRANTY; without even the implied warranty of              #
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the               #
# GNU General Public License for more details.                                #
#                                                                             #
# You should have received a copy of the GNU General Public License           #
# along with this program.  If not, see <http://www.gnu.org/licenses/>.       #
#                                                                             #
###############################################################################

import logging
import os
import parted
import stat

from . import step
from . import util
from .errors import *
from .i18n import _

# Setup logging
log = logging.getLogger("bricklayer.disk")

DEFAULT_FILESYSTEM = "btrfs"

# Check if parted.PARTITION_ESP is set
try:
	parted.PARTITION_ESP
except AttributeError:
	parted.PARTITION_ESP = 18

class Disks(object):
	"""
		Disks abstraction class
	"""
	def __init__(self, bricklayer):
		self.bricklayer = bricklayer

		# Disks
		self.disks = []

	def scan(self):
		"""
			Scans for all disks
		"""
		# Don't scan for disks if already done
		if self.disks:
			return

		log.debug("Scanning for disks...")

		for device in parted.getAllDevices():
			disk = Disk(self.bricklayer, device)

			# Skip whatever isn't suitable
			if not disk.supported:
				continue

			self.disks.append(disk)

		# Sort them alphabetically
		self.disks.sort()

	def add_disk(self, path, selected=False):
		"""
			Adds the disk at path
		"""
		# Check if the disk is already on the list
		for disk in self.disks:
			if disk.path == path:
				return disk

		st = os.stat(path)

		# Setup regular files as loop devices
		if stat.S_ISREG(st.st_mode):
			path = self._losetup(path)

		# Create a Disk object
		disk = Disk(self.bricklayer, path)
		self.disks.append(disk)

		# Select this disk
		if selected:
			disk.selected = True

		return disk

	def _losetup(self, path):
		# Find a free loop device
		device = self.bricklayer.command(["losetup", "-f"])
		device = device.rstrip()

		# Connect the image to the loop device
		self.bricklayer.command(["losetup", device, path])

		# Return the name of the loop device
		return device

	@property
	def supported(self):
		"""
			The supported disks
		"""
		return [disk for disk in self.disks if disk.supported]

	@property
	def selected(self):
		"""
			The selected disks
		"""
		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()

	def _find_partition(self, name):
		"""
			Returns the partition with name
		"""
		for disk in self.selected:
			for partition in disk.partitions:
				if partition.name == name:
					return partition

	def mount(self):
		"""
			Mounts all filesystems
		"""
		# Find root partition
		partition = self._find_partition("ROOT")
		if not partition:
			FileNotFoundError("Could not find root partition")

		# Mount the root partition
		self._mount(partition.path, self.bricklayer.root)

		# Find ESP partition
		partition = self._find_partition("ESP")
		if partition:
			self._mount(partition.path, os.path.join(self.bricklayer.root, "boot/efi"))

	def _mount(self, source, target):
		"""
			Mounts source to target
		"""
		# Make sure the target exists
		os.makedirs(target, exist_ok=True)

		# Call mount(8)
		self.bricklayer.command(["mount", source, target])

	def umount(self):
		"""
			Umounts all filesystems
		"""
		# Umount everything mounted in root
		self.bricklayer.command(["umount", "-Rv", self.bricklayer.root], error_ok=True)

	def tear_down(self):
		"""
			Shuts down any storage
		"""
		self.umount()

		for disk in self.disks:
			disk.tear_down()


class Disk(object):
	def __init__(self, bricklayer, device):
		self.bricklayer = bricklayer

		# The parted device
		if not isinstance(device, parted.Device):
			device = parted.Device(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.path)

	def __str__(self):
		return "%s (%s)" % (self.model, util.format_size(self.size))

	def __hash__(self):
		return hash(self.path)

	def __eq__(self, other):
		if isinstance(other, self.__class__):
			return self.device == other.device

		return NotImplemented

	def __lt__(self, other):
		if isinstance(other, self.__class__):
			return self.model < other.model or self.path < other.path

		return NotImplemented

	@property
	def supported(self):
		"""
			Is this device supported?
		"""
		# Skip any device-mapper devices
		if self.device.type == parted.DEVICE_DM:
			return False

		# Skip any CD/DVD drives
		if self.device.path.startswith("/dev/sr"):
			return False

		# Skip MDRAID devices
		if self.device.path.startswith("/dev/md"):
			return False

		# We do not support read-only devices
		if self.device.readOnly:
			return False

		# Ignore any busy devices
		if self.device.busy:
			return False

		return True

	@property
	def path(self):
		return self.device.path

	@property
	def model(self):
		return self.device.model or _("Unknown Model")

	@property
	def size(self):
		return self.device.length * self.device.sectorSize

	@property
	def partitions(self):
		"""
			Returns a list of all partitions on this device
		"""
		return [Partition(self.bricklayer, p) for p in self.parted.partitions]

	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
		if any((bl.requires_bootldr_partition for bl in self.bricklayer.bootloaders)):
			self._add_partition("BOOTLDR", DEFAULT_FILESYSTEM, length=1024**2,
				flags=[parted.PARTITION_BIOS_GRUB])

		# Create an EFI-partition of exactly 32 MiB
		if any((bl.requires_efi_partition for bl in self.bricklayer.bootloaders)):
			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
		"""
		# Destroy any existing partition table
		self.device.clobber()

		# Write the new partition table
		self.parted.commit()

		# For loop devices, we have to manually create the partition mappings
		if self.path.startswith("/dev/loop"):
			self.bricklayer.command(["kpartx", "-av", self.path])

	def tear_down(self):
		"""
			Shuts down this disk
		"""
		# Unmap partitions on loop devices
		if self.path.startswith("/dev/loop"):
			self.bricklayer.command(["kpartx", "-dv", self.path])

			# Free the loop device
			self.bricklayer.command(["losetup", "-d", self.path])


class Partition(object):
	def __init__(self, bricklayer, parted):
		self.bricklayer = bricklayer

		# The parted device
		self.parted = parted

	def __repr__(self):
		return "<%s %s>" % (self.__class__.__name__, self.name or self.path)

	@property
	def name(self):
		return self.parted.name

	@property
	def path(self):
		# Map path for loop devices
		if self.parted.path.startswith("/dev/loop"):
			return self.parted.path.replace("/dev/loop", "/dev/mapper/loop")

		return self.parted.path

	def wipe(self):
		"""
			Wipes the entire partition (i.e. writes zeroes)
		"""
		log.info("Wiping %s (%s)..." % (self.name, self.path))

		zero = bytearray(1024)

		with open(self.path, "wb") as f:
			f.write(zero)

	def format(self):
		"""
			Formats the filesystem
		"""
		# Wipe BIOS_GRUB partitions instead of formatting them
		if self.parted.getFlag(parted.PARTITION_BIOS_GRUB):
			return self.wipe()

		# Fetch file-system type
		filesystem = self.parted.fileSystem.type

		log.info("Formatting %s (%s) with %s..." % (self.name, self.path, filesystem))

		if filesystem == "fat32":
			command = ["mkfs.vfat", self.path]
		elif filesystem == "linux-swap(v1)":
			command = ["mkswap", "-v1", self.path]
		else:
			command = ["mkfs.%s" % filesystem, "-f", self.path]

		# Run command
		self.bricklayer.command(command)


class UnattendedSelectDisk(step.UnattendedStep):
	"""
		Scans for any disks
	"""
	def run(self):
		# Nothing to do if disks have already been selected on the CLI
		if self.bricklayer.disks.selected:
			return

		# Scan for disks
		self.bricklayer.disks.scan()

		# End here if we could not find any disks
		if not self.bricklayer.disks.supported:
			self.tui.error(
				_("No Disks Found"),
				_("No supported disks were found")
			)

			raise InstallAbortedError("No disks found")

		# Automatically select the first disk
		for disk in self.bricklayer.disks.supported:
			disk.selected = True
			break


class SelectDisk(step.InteractiveStep):
	"""
		Ask the user which disk(s) to use for the installation process
	"""
	def run(self):
		# Scan for disks
		self.bricklayer.disks.scan()

		# Create a dictionary with all disks
		disks = { disk : "%s" % disk for disk in self.bricklayer.disks.supported }

		# Show an error if no suitable disks were found
		if not disks:
			self.tui.error(
				_("No Disks Found"),
				_("No supported disks were found")
			)

			raise InstallAbortedError("No disks found")

		# Get the current selection
		selection = [disk for disk in disks if disk.selected]

		while True:
			# Select disks
			selection = self.tui.select(
				_("Disk Selection"),
				_("Please select all disks for installation"),
				disks, default=selection, multi=True, width=60,
			)

			# Is at least one disk selected?
			if not selection:
				self.tui.error(
					_("No Disk Selected"),
					_("Please select a disk to continue the installation"),
					buttons=[_("Back")],
				)
				continue

			# Apply selection
			for disk in disks:
				disk.selected = disk in selection

			break


class CalculatePartitionLayout(step.Step):
	"""
		Calculates the partition layout
	"""
	def run(self):
		# 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):
		log.debug("Creating partitions")

		with self.tui.progress(
			_("Creating Partition Layout"),
			_("Creating partition layout..."),
		):
			for disk in self.bricklayer.disks.selected:
				disk.commit()


class CreateFilesystems(step.Step):
	"""
		Formats all newly created partitions
	"""
	def run(self):
		for disk in self.bricklayer.disks.selected:
			for partition in disk.partitions:
				with self.tui.progress(
					_("Creating Filesystems"),
					_("Formatting partition \"%s\"...") % (partition.name or partition.path)
				):
					partition.format()


class MountFilesystems(step.Step):
	"""
		Mount all filesystems
	"""
	def run(self):
		with self.tui.progress(
			_("Mounting Filesystems"),
			_("Mounting filesystems..."),
		):
			self.bricklayer.disks.mount()


class UmountFilesystems(step.Step):
	"""
		Umount all filesystems
	"""
	def run(self):
		with self.tui.progress(
			_("Umounting Filesystems"),
			_("Umounting filesystems..."),
		):
			self.bricklayer.disks.umount()