image: T4516: add raid-1 install support

(cherry picked from commit e036f783bc85e4d2bad5f5cbfd688a03a352223e)

4 files changed, 196 insertions, 27 deletions

diff --git a/python/vyos/system/disk.py b/python/vyos/system/disk.py
index 882b4eb39..49e6b5c5e 100644
--- a/python/vyos/system/disk.py
+++ b/python/vyos/system/disk.py
@@ -15,12 +15,20 @@
 
 from json import loads as json_loads
 from os import sync
+from dataclasses import dataclass
 
 from psutil import disk_partitions
 
 from vyos.utils.process import run, cmd
 
 
+@dataclass
+class DiskDetails:
+    """Disk details"""
+    name: str
+    partition: dict[str, str]
+
+
 def disk_cleanup(drive_path: str) -> None:
     """Clean up disk partition table (MBR and GPT)
     Zeroize primary and secondary headers - first and last 17408 bytes
@@ -67,6 +75,62 @@ def parttable_create(drive_path: str, root_size: int) -> None:
     sync()
     run(f'partprobe {drive_path}')
 
+    partitions: list[str] = partition_list(drive_path)
+
+    disk: DiskDetails = DiskDetails(
+        name = drive_path,
+        partition = {
+            'efi': next(x for x in partitions if x.endswith('2')),
+            'root': next(x for x in partitions if x.endswith('3'))
+        }
+    )
+
+    return disk
+
+
+def partition_list(drive_path: str) -> list[str]:
+    """Get a list of partitions on a drive
+
+    Args:
+        drive_path (str): path to a drive
+
+    Returns:
+        list[str]: a list of partition paths
+    """
+    lsblk: str = cmd(f'lsblk -Jp {drive_path}')
+    drive_info: dict = json_loads(lsblk)
+    device: list = drive_info.get('blockdevices')
+    children: list[str] = device[0].get('children', []) if device else []
+    partitions: list[str] = [child.get('name') for child in children]
+    return partitions
+
+
+def partition_parent(partition_path: str) -> str:
+    """Get a parent device for a partition
+
+    Args:
+        partition (str): path to a partition
+
+    Returns:
+        str: path to a parent device
+    """
+    parent: str = cmd(f'lsblk -ndpo pkname {partition_path}')
+    return parent
+
+
+def from_partition(partition_path: str) -> DiskDetails:
+    drive_path: str = partition_parent(partition_path)
+    partitions: list[str] = partition_list(drive_path)
+
+    disk: DiskDetails = DiskDetails(
+        name = drive_path,
+        partition = {
+            'efi': next(x for x in partitions if x.endswith('2')),
+            'root': next(x for x in partitions if x.endswith('3'))
+        }
+    )
+
+    return disk
 
 def filesystem_create(partition: str, fstype: str) -> None:
     """Create a filesystem on a partition
@@ -138,25 +202,6 @@ def find_device(mountpoint: str) -> str:
     return ''
 
 
-def raid_create(raid_name: str,
-                raid_members: list[str],
-                raid_level: str = 'raid1') -> None:
-    """Create a RAID array
-
-    Args:
-        raid_name (str): a name of array (data, backup, test, etc.)
-        raid_members (list[str]): a list of array members
-        raid_level (str, optional): an array level. Defaults to 'raid1'.
-    """
-    raid_devices_num: int = len(raid_members)
-    raid_members_str: str = ' '.join(raid_members)
-    command: str = f'mdadm --create /dev/md/{raid_name} --metadata=1.2 \
-        --raid-devices={raid_devices_num} --level={raid_level} \
-        {raid_members_str}'
-
-    run(command)
-
-
 def disks_size() -> dict[str, int]:
     """Get a dictionary with physical disks and their sizes
 
diff --git a/python/vyos/system/grub.py b/python/vyos/system/grub.py
index 9ac205c03..0ac16af9a 100644
--- a/python/vyos/system/grub.py
+++ b/python/vyos/system/grub.py
@@ -19,7 +19,7 @@ from typing import Union
 from uuid import uuid5, NAMESPACE_URL, UUID
 
 from vyos.template import render
-from vyos.utils.process import run, cmd
+from vyos.utils.process import cmd
 from vyos.system import disk
 
 # Define variables
@@ -49,7 +49,7 @@ REGEX_GRUB_MODULES: str = r'^insmod (?P<module_name>.+)$'
 REGEX_KERNEL_CMDLINE: str = r'^BOOT_IMAGE=/(?P<boot_type>boot|live)/((?P<image_version>.+)/)?vmlinuz.*$'
 
 
-def install(drive_path: str, boot_dir: str, efi_dir: str) -> None:
+def install(drive_path: str, boot_dir: str, efi_dir: str, id: str = 'VyOS') -> None:
     """Install GRUB for both BIOS and EFI modes (hybrid boot)
 
     Args:
@@ -62,11 +62,11 @@ def install(drive_path: str, boot_dir: str, efi_dir: str) -> None:
             {drive_path} --force',
         f'grub-install --no-floppy --recheck --target=x86_64-efi \
             --force-extra-removable --boot-directory={boot_dir} \
-            --efi-directory={efi_dir} --bootloader-id="VyOS" \
+            --efi-directory={efi_dir} --bootloader-id="{id}" \
             --no-uefi-secure-boot'
     ]
     for command in commands:
-        run(command)
+        cmd(command)
 
 
 def gen_version_uuid(version_name: str) -> str:
@@ -294,7 +294,7 @@ def set_default(version_name: str, root_dir: str = '') -> None:
     vars_write(vars_file, vars_current)
 
 
-def common_write(root_dir: str = '') -> None:
+def common_write(root_dir: str = '', grub_common: dict[str, str] = {}) -> None:
     """Write common GRUB configuration file (overwrite everything)
 
     Args:
@@ -304,7 +304,7 @@ def common_write(root_dir: str = '') -> None:
     if not root_dir:
         root_dir = disk.find_persistence()
     common_config = f'{root_dir}/{CFG_VYOS_COMMON}'
-    render(common_config, TMPL_GRUB_COMMON, {})
+    render(common_config, TMPL_GRUB_COMMON, grub_common)
 
 
 def create_structure(root_dir: str = '') -> None:
@@ -335,3 +335,6 @@ def set_console_type(console_type: str, root_dir: str = '') -> None:
     vars_current: dict[str, str] = vars_read(vars_file)
     vars_current['console_type'] = str(console_type)
     vars_write(vars_file, vars_current)
+
+def set_raid(root_dir: str = '') -> None:
+    pass
diff --git a/python/vyos/system/raid.py b/python/vyos/system/raid.py
new file mode 100644
index 000000000..13b99fa69
--- /dev/null
+++ b/python/vyos/system/raid.py
@@ -0,0 +1,115 @@
+# Copyright 2023 VyOS maintainers and contributors <maintainers@vyos.io>
+#
+# This library is free software; you can redistribute it and/or
+# modify it under the terms of the GNU Lesser General Public
+# License as published by the Free Software Foundation; either
+# version 2.1 of the License, or (at your option) any later version.
+#
+# This library 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
+# Lesser General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public License
+# along with this library.  If not, see <http://www.gnu.org/licenses/>.
+
+"""RAID related functions"""
+
+from pathlib import Path
+from shutil import copy
+from dataclasses import dataclass
+
+from vyos.utils.process import cmd
+from vyos.system import disk
+
+
+@dataclass
+class RaidDetails:
+    """RAID type"""
+    name: str
+    level: str
+    members: list[str]
+    disks: list[disk.DiskDetails]
+
+
+def raid_create(raid_members: list[str],
+                raid_name: str = 'md0',
+                raid_level: str = 'raid1') -> None:
+    """Create a RAID array
+
+    Args:
+        raid_name (str): a name of array (data, backup, test, etc.)
+        raid_members (list[str]): a list of array members
+        raid_level (str, optional): an array level. Defaults to 'raid1'.
+    """
+    raid_devices_num: int = len(raid_members)
+    raid_members_str: str = ' '.join(raid_members)
+    if Path('/sys/firmware/efi').exists():
+        for part in raid_members:
+            drive: str = disk.partition_parent(part)
+            command: str = f'sgdisk --typecode=3:A19D880F-05FC-4D3B-A006-743F0F84911E {drive}'
+            cmd(command)
+    else:
+        for part in raid_members:
+            drive: str = disk.partition_parent(part)
+            command: str = f'sgdisk --typecode=3:A19D880F-05FC-4D3B-A006-743F0F84911E {drive}'
+            cmd(command)
+    for part in raid_members:
+        command: str = f'mdadm --zero-superblock {part}'
+        cmd(command)
+    command: str = f'mdadm --create /dev/{raid_name} -R --metadata=1.0 \
+        --raid-devices={raid_devices_num} --level={raid_level} \
+        {raid_members_str}'
+
+    cmd(command)
+
+    raid = RaidDetails(
+        name = f'/dev/{raid_name}',
+        level = raid_level,
+        members = raid_members,
+        disks = [disk.from_partition(m) for m in raid_members]
+    )
+
+    return raid
+
+def update_initramfs() -> None:
+    """Update initramfs"""
+    mdadm_script = '/etc/initramfs-tools/scripts/local-top/mdadm'
+    copy('/usr/share/initramfs-tools/scripts/local-block/mdadm', mdadm_script)
+    p = Path(mdadm_script)
+    p.write_text(p.read_text().replace('$((COUNT + 1))', '20'))
+    command: str = 'update-initramfs -u'
+    cmd(command)
+
+def update_default(target_dir: str) -> None:
+    """Update /etc/default/mdadm to start MD monitoring daemon at boot
+    """
+    source_mdadm_config = '/etc/default/mdadm'
+    target_mdadm_config = Path(target_dir).joinpath('/etc/default/mdadm')
+    target_mdadm_config_dir = Path(target_mdadm_config).parent
+    Path.mkdir(target_mdadm_config_dir, parents=True, exist_ok=True)
+    s = Path(source_mdadm_config).read_text().replace('START_DAEMON=false',
+                                                      'START_DAEMON=true')
+    Path(target_mdadm_config).write_text(s)
+
+def get_uuid(device: str) -> str:
+    """Get UUID of a device"""
+    command: str = f'tune2fs -l {device}'
+    l = cmd(command).splitlines()
+    uuid = next((x for x in l if x.startswith('Filesystem UUID')), '')
+    return uuid.split(':')[1].strip() if uuid else ''
+
+def get_uuids(raid_details: RaidDetails) -> tuple[str]:
+    """Get UUIDs of RAID members
+
+    Args:
+        raid_name (str): a name of array (data, backup, test, etc.)
+
+    Returns:
+        tuple[str]: root_disk uuid, root_md uuid
+    """
+    raid_name: str = raid_details.name
+    root_partition: str = raid_details.members[0]
+    uuid_root_disk: str = get_uuid(root_partition)
+    uuid_root_md: str = get_uuid(raid_name)
+    return uuid_root_disk, uuid_root_md
diff --git a/python/vyos/utils/io.py b/python/vyos/utils/io.py
index e34a1ba32..74099b502 100644
--- a/python/vyos/utils/io.py
+++ b/python/vyos/utils/io.py
@@ -13,6 +13,8 @@
 # You should have received a copy of the GNU Lesser General Public
 # License along with this library.  If not, see <http://www.gnu.org/licenses/>.
 
+from typing import Callable
+
 def print_error(str='', end='\n'):
     """
     Print `str` to stderr, terminated with `end`.
@@ -73,7 +75,8 @@ def is_dumb_terminal():
     import os
     return os.getenv('TERM') in ['vt100', 'dumb']
 
-def select_entry(l: list, list_msg: str = '', prompt_msg: str = '') -> str:
+def select_entry(l: list, list_msg: str = '', prompt_msg: str = '',
+                 list_format: Callable = None,) -> str:
     """Select an entry from a list
 
     Args:
@@ -87,7 +90,10 @@ def select_entry(l: list, list_msg: str = '', prompt_msg: str = '') -> str:
     en = list(enumerate(l, 1))
     print(list_msg)
     for i, e in en:
-        print(f'\t{i}: {e}')
+        if list_format:
+            print(f'\t{i}: {list_format(e)}')
+        else:
+            print(f'\t{i}: {e}')
     select = ask_input(prompt_msg, numeric_only=True,
                        valid_responses=range(1, len(l)+1))
     return next(filter(lambda x: x[0] == select, en))[1]