Add support for Btrfs snapshots (#3500)

* Add btrfs snapshot support

* Update

* Update

* Update

Author: svartkanin

archinstall/lib/disk/disk_menu.py

@@ -1,9 +1,19 @@
 from dataclasses import dataclass
 from typing import override
 
-from archinstall.lib.models.device_model import DiskLayoutConfiguration, DiskLayoutType, LvmConfiguration
+from archinstall.lib.models.device_model import (
+	BtrfsOptions,
+	DiskLayoutConfiguration,
+	DiskLayoutType,
+	LvmConfiguration,
+	SnapshotConfig,
+	SnapshotType,
+)
 from archinstall.lib.translationhandler import tr
+from archinstall.tui.curses_menu import SelectMenu
 from archinstall.tui.menu_item import MenuItem, MenuItemGroup
+from archinstall.tui.result import ResultType
+from archinstall.tui.types import Alignment, FrameProperties
 
 from ..interactions.disk_conf import select_disk_config, select_lvm_config
 from ..menu.abstract_menu import AbstractSubMenu
@@ -14,16 +24,24 @@
 class DiskMenuConfig:
 	disk_config: DiskLayoutConfiguration | None
 	lvm_config: LvmConfiguration | None
+	btrfs_snapshot_config: SnapshotConfig | None
 
 
 class DiskLayoutConfigurationMenu(AbstractSubMenu[DiskLayoutConfiguration]):
 	def __init__(self, disk_layout_config: DiskLayoutConfiguration | None):
 		if not disk_layout_config:
-			self._disk_menu_config = DiskMenuConfig(disk_config=None, lvm_config=None)
+			self._disk_menu_config = DiskMenuConfig(
+				disk_config=None,
+				lvm_config=None,
+				btrfs_snapshot_config=None,
+			)
 		else:
+			snapshot_config = disk_layout_config.btrfs_options.snapshot_config if disk_layout_config.btrfs_options else None
+
 			self._disk_menu_config = DiskMenuConfig(
 				disk_config=disk_layout_config,
 				lvm_config=disk_layout_config.lvm_config,
+				btrfs_snapshot_config=snapshot_config,
 			)
 
 		menu_optioons = self._define_menu_options()
@@ -52,6 +70,14 @@ def _define_menu_options(self) -> list[MenuItem]:
 				dependencies=[self._check_dep_lvm],
 				key='lvm_config',
 			),
+			MenuItem(
+				text='Btrfs snapshots',
+				action=self._select_btrfs_snapshots,
+				value=self._disk_menu_config.btrfs_snapshot_config,
+				preview_action=self._prev_btrfs_snapshots,
+				dependencies=[self._check_dep_btrfs],
+				key='btrfs_snapshot_config',
+			),
 		]
 
 	@override
@@ -60,6 +86,7 @@ def run(self) -> DiskLayoutConfiguration | None:
 
 		if self._disk_menu_config.disk_config:
 			self._disk_menu_config.disk_config.lvm_config = self._disk_menu_config.lvm_config
+			self._disk_menu_config.disk_config.btrfs_options = BtrfsOptions(snapshot_config=self._disk_menu_config.btrfs_snapshot_config)
 			return self._disk_menu_config.disk_config
 
 		return None
@@ -72,10 +99,15 @@ def _check_dep_lvm(self) -> bool:
 
 		return False
 
-	def _select_disk_layout_config(
-		self,
-		preset: DiskLayoutConfiguration | None,
-	) -> DiskLayoutConfiguration | None:
+	def _check_dep_btrfs(self) -> bool:
+		disk_layout_conf: DiskLayoutConfiguration | None = self._menu_item_group.find_by_key('disk_config').value
+
+		if disk_layout_conf:
+			return disk_layout_conf.is_default_btrfs()
+
+		return False
+
+	def _select_disk_layout_config(self, preset: DiskLayoutConfiguration | None) -> DiskLayoutConfiguration | None:
 		disk_config = select_disk_config(preset)
 
 		if disk_config != preset:
@@ -91,6 +123,38 @@ def _select_lvm_config(self, preset: LvmConfiguration | None) -> LvmConfiguratio
 
 		return preset
 
+	def _select_btrfs_snapshots(self, preset: SnapshotConfig | None) -> SnapshotConfig | None:
+		preset_type = preset.snapshot_type if preset else None
+
+		group = MenuItemGroup.from_enum(
+			SnapshotType,
+			sort_items=True,
+			preset=preset_type,
+		)
+
+		result = SelectMenu[SnapshotType](
+			group,
+			allow_reset=True,
+			allow_skip=True,
+			frame=FrameProperties.min(tr('Snapshot type')),
+			alignment=Alignment.CENTER,
+		).run()
+
+		snapshot_type: SnapshotType | None = None
+
+		match result.type_:
+			case ResultType.Skip:
+				return preset
+			case ResultType.Reset:
+				return None
+			case ResultType.Selection:
+				snapshot_type = result.get_value()
+
+		if not snapshot_type:
+			return None
+
+		return SnapshotConfig(snapshot_type=snapshot_type)
+
 	def _prev_disk_layouts(self, item: MenuItem) -> str | None:
 		if not item.value:
 			return None
@@ -146,3 +210,10 @@ def _prev_lvm_config(self, item: MenuItem) -> str | None:
 			return output
 
 		return None
+
+	def _prev_btrfs_snapshots(self, item: MenuItem) -> str | None:
+		if not item.value:
+			return None
+
+		snapshot_config: SnapshotConfig = item.value
+		return tr('Snapshot type: {}').format(snapshot_config.snapshot_type.value)

archinstall/lib/global_menu.py

@@ -335,6 +335,11 @@ def _prev_disk_config(self, item: MenuItem) -> str | None:
 			if disk_layout_conf.lvm_config:
 				output += '{}: {}'.format(tr('LVM configuration type'), disk_layout_conf.lvm_config.config_type.display_msg())
 
+			if disk_layout_conf.btrfs_options:
+				btrfs_options = disk_layout_conf.btrfs_options
+				if btrfs_options.snapshot_config:
+					output += tr('Btrfs snapshot type: {}').format(btrfs_options.snapshot_config.snapshot_type.value)
+
 			return output
 
 		return None

archinstall/lib/installer.py

@@ -25,6 +25,7 @@
 	PartitionModification,
 	SectorSize,
 	Size,
+	SnapshotType,
 	SubvolumeModification,
 	Unit,
 )
@@ -909,6 +910,51 @@ def minimal_installation(
 			if hasattr(plugin, 'on_install'):
 				plugin.on_install(self)
 
+	def setup_btrfs_snapshot(
+		self,
+		snapshot_type: SnapshotType,
+		bootloader: Bootloader | None = None,
+	) -> None:
+		if snapshot_type == SnapshotType.Snapper:
+			debug('Setting up Btrfs snapper')
+			self.pacman.strap('snapper')
+
+			snapper: dict[str, str] = {
+				'root': '/',
+				'home': '/home',
+			}
+
+			for config_name, mountpoint in snapper.items():
+				command = [
+					'arch-chroot',
+					str(self.target),
+					'snapper',
+					'--no-dbus',
+					'-c',
+					config_name,
+					'create-config',
+					mountpoint,
+				]
+
+				try:
+					SysCommand(command, peek_output=True)
+				except SysCallError as err:
+					raise DiskError(f'Could not setup Btrfs snapper: {err}')
+
+			self.enable_service('snapper-timeline.timer')
+			self.enable_service('snapper-cleanup.timer')
+		elif snapshot_type == SnapshotType.Timeshift:
+			debug('Setting up Btrfs timeshift')
+
+			self.pacman.strap('cronie')
+			self.pacman.strap('timeshift')
+
+			self.enable_service('cronie.service')
+
+			if bootloader and bootloader == Bootloader.Grub:
+				self.pacman.strap('grub-btrfs')
+				self.enable_service('grub-btrfs.service')
+
 	def setup_swap(self, kind: str = 'zram') -> None:
 		if kind == 'zram':
 			info('Setting up swap on zram')

archinstall/lib/models/device_model.py

@@ -40,13 +40,15 @@ class _DiskLayoutConfigurationSerialization(TypedDict):
 	device_modifications: NotRequired[list[_DeviceModificationSerialization]]
 	lvm_config: NotRequired[_LvmConfigurationSerialization]
 	mountpoint: NotRequired[str]
+	btrfs_options: NotRequired[_BtrfsOptionsSerialization]
 
 
 @dataclass
 class DiskLayoutConfiguration:
 	config_type: DiskLayoutType
 	device_modifications: list[DeviceModification] = field(default_factory=list)
 	lvm_config: LvmConfiguration | None = None
+	btrfs_options: BtrfsOptions | None = None
 
 	# used for pre-mounted config
 	mountpoint: Path | None = None
@@ -66,6 +68,9 @@ def json(self) -> _DiskLayoutConfigurationSerialization:
 			if self.lvm_config:
 				config['lvm_config'] = self.lvm_config.json()
 
+			if self.btrfs_options:
+				config['btrfs_options'] = self.btrfs_options.json()
+
 			return config
 
 	@classmethod
@@ -174,8 +179,22 @@ def parse_arg(cls, disk_config: _DiskLayoutConfigurationSerialization) -> DiskLa
 		if (lvm_arg := disk_config.get('lvm_config', None)) is not None:
 			config.lvm_config = LvmConfiguration.parse_arg(lvm_arg, config)
 
+		if config.is_default_btrfs():
+			if (btrfs_arg := disk_config.get('btrfs_options', None)) is not None:
+				config.btrfs_options = BtrfsOptions.parse_arg(btrfs_arg)
+
 		return config
 
+	def is_default_btrfs(self) -> bool:
+		if self.config_type == DiskLayoutType.Default:
+			for mod in self.device_modifications:
+				for part in mod.partitions:
+					if part.is_create_or_modify():
+						if part.fs_type == FilesystemType.Btrfs:
+							return True
+
+		return False
+
 
 class PartitionTable(Enum):
 	GPT = 'gpt'
@@ -1309,13 +1328,46 @@ def get_root_volume(self) -> LvmVolume | None:
 		return None
 
 
-# def get_lv_crypt_uuid(self, lv: LvmVolume, encryption: EncryptionType) -> str:
-# 	"""
-# 	Find the LUKS superblock UUID for the device that
-# 	contains the given logical volume
-# 	"""
-# 	for vg in self.vol_groups:
-# 		if vg.contains_lv(lv):
+class _BtrfsOptionsSerialization(TypedDict):
+	snapshot_config: _SnapshotConfigSerialization | None
+
+
+class _SnapshotConfigSerialization(TypedDict):
+	type: str
+
+
+class SnapshotType(Enum):
+	Snapper = 'Snapper'
+	Timeshift = 'Timeshift'
+
+
+@dataclass
+class SnapshotConfig:
+	snapshot_type: SnapshotType
+
+	def json(self) -> _SnapshotConfigSerialization:
+		return {'type': self.snapshot_type.value}
+
+	@staticmethod
+	def parse_args(args: _SnapshotConfigSerialization) -> SnapshotConfig | None:
+		return SnapshotConfig(SnapshotType(args['type']))
+
+
+@dataclass
+class BtrfsOptions:
+	snapshot_config: SnapshotConfig | None
+
+	def json(self) -> _BtrfsOptionsSerialization:
+		return {'snapshot_config': self.snapshot_config.json() if self.snapshot_config else None}
+
+	@staticmethod
+	def parse_arg(arg: _BtrfsOptionsSerialization) -> BtrfsOptions | None:
+		snapshot_args = arg.get('snapshot_config')
+		if snapshot_args:
+			snapshot_config = SnapshotConfig.parse_args(snapshot_args)
+			return BtrfsOptions(snapshot_config)
+
+		return None
 
 
 class _DeviceModificationSerialization(TypedDict):

archinstall/scripts/guided.py

@@ -141,6 +141,13 @@ def perform_installation(mountpoint: Path) -> None:
 		if servies := config.services:
 			installation.enable_service(servies)
 
+		if disk_config.is_default_btrfs():
+			btrfs_options = disk_config.btrfs_options
+			snapshot_config = btrfs_options.snapshot_config if btrfs_options else None
+			snapshot_type = snapshot_config.snapshot_type if snapshot_config else None
+			if snapshot_type:
+				installation.setup_btrfs_snapshot(snapshot_type, config.bootloader)
+
 		# If the user provided custom commands to be run post-installation, execute them now.
 		if cc := config.custom_commands:
 			run_custom_user_commands(cc, installation)

archinstall/tui/menu_item.py

@@ -2,6 +2,7 @@
 
 from collections.abc import Callable
 from dataclasses import dataclass, field
+from enum import Enum
 from functools import cached_property
 from typing import Any, ClassVar
 
@@ -119,6 +120,20 @@ def yes_no() -> 'MenuItemGroup':
 			sort_items=True,
 		)
 
+	@staticmethod
+	def from_enum(
+		enum_cls: type[Enum],
+		sort_items: bool = False,
+		preset: Enum | None = None,
+	) -> 'MenuItemGroup':
+		items = [MenuItem(elem.value, value=elem) for elem in enum_cls]
+		group = MenuItemGroup(items, sort_items=False)
+
+		if preset is not None:
+			group.set_selected_by_value(preset)
+
+		return group
+
 	def set_preview_for_all(self, action: Callable[[Any], str | None]) -> None:
 		for item in self.items:
 			item.preview_action = action

tests/conftest.py

@@ -8,6 +8,11 @@ def config_fixture() -> Path:
 	return Path(__file__).parent / 'data' / 'test_config.json'
 
 
+@pytest.fixture(scope='session')
+def btrfs_config_fixture() -> Path:
+	return Path(__file__).parent / 'data' / 'test_config_btrfs.json'
+
+
 @pytest.fixture(scope='session')
 def creds_fixture() -> Path:
 	return Path(__file__).parent / 'data' / 'test_creds.json'

tests/data/test_config.json

@@ -10,6 +10,11 @@
   ],
   "disk_config": {
     "config_type": "default_layout",
+    "btrfs_options": {
+      "snapshot_config": {
+        "type": "Snapper"
+      }
+    },
     "device_modifications": [
       {
         "device": "/dev/sda",