# Copyright 2023 VyOS maintainers and contributors # # 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 . from json import loads as json_loads from os import sync from dataclasses import dataclass from time import sleep 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) Remove partition and device signatures. Zeroize primary and secondary headers - first and last 17408 bytes (512 bytes * 34 LBA) on a drive Args: drive_path (str): path to a drive that needs to be cleaned """ partitions: list[str] = partition_list(drive_path) for partition in partitions: run(f'wipefs -af {partition}') run(f'wipefs -af {drive_path}') run(f'sgdisk -Z {drive_path}') def find_persistence() -> str: """Find a mountpoint for persistence storage Returns: str: Path where 'persistance' pertition is mounted, Empty if not found """ mounted_partitions = disk_partitions() for partition in mounted_partitions: if partition.mountpoint.endswith('/persistence'): return partition.mountpoint return '' def parttable_create(drive_path: str, root_size: int) -> None: """Create a hybrid MBR/GPT partition table 0-2047 first sectors are free 2048-4095 sectors - BIOS Boot Partition 4096 + 256 MB - EFI system partition Everything else till the end of a drive - Linux partition Args: drive_path (str): path to a drive """ if not root_size: root_size_text: str = '+100%' else: root_size_text: str = str(root_size) command = f'sgdisk -a1 -n1:2048:4095 -t1:EF02 -n2:4096:+256M -t2:EF00 \ -n3:0:+{root_size_text}K -t3:8300 {drive_path}' run(command) # update partitons in kernel sync() run(f'partx -u {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 Args: partition (str): path to a partition (for example: '/dev/sda1') fstype (str): filesystem type ('efi' or 'ext4') """ if fstype == 'efi': command = 'mkfs -t fat -n EFI' run(f'{command} {partition}') if fstype == 'ext4': command = 'mkfs -t ext4 -L persistence' run(f'{command} {partition}') def partition_mount(partition: str, path: str, fsype: str = '', overlay_params: dict[str, str] = {}) -> bool: """Mount a partition into a path Args: partition (str): path to a partition (for example: '/dev/sda1') path (str): a path where to mount fsype (str): optionally, set fstype ('squashfs', 'overlay', 'iso9660') overlay_params (dict): optionally, set overlay parameters. Defaults to None. Returns: bool: True on success """ if fsype in ['squashfs', 'iso9660']: command: str = f'mount -o loop,ro -t {fsype} {partition} {path}' if fsype == 'overlay' and overlay_params: command: str = f'mount -t overlay -o noatime,\ upperdir={overlay_params["upperdir"]},\ lowerdir={overlay_params["lowerdir"]},\ workdir={overlay_params["workdir"]} overlay {path}' else: command = f'mount {partition} {path}' rc = run(command) if rc == 0: return True return False def partition_umount(partition: str = '', path: str = '') -> None: """Umount a partition by a partition name or a path Args: partition (str): path to a partition (for example: '/dev/sda1') path (str): a path where a partition is mounted """ if partition: command = f'umount {partition}' run(command) if path: command = f'umount {path}' run(command) def find_device(mountpoint: str) -> str: """Find a device by mountpoint Returns: str: Path to device, Empty if not found """ mounted_partitions = disk_partitions(all=True) for partition in mounted_partitions: if partition.mountpoint == mountpoint: return partition.mountpoint return '' def wait_for_umount(mountpoint: str = '') -> None: """Wait (within reason) for umount to complete """ i = 0 while find_device(mountpoint): i += 1 if i == 5: print(f'Warning: {mountpoint} still mounted') break sleep(1) def disks_size() -> dict[str, int]: """Get a dictionary with physical disks and their sizes Returns: dict[str, int]: a dictionary with name: size mapping """ disks_size: dict[str, int] = {} lsblk: str = cmd('lsblk -Jbp') blk_list = json_loads(lsblk) for device in blk_list.get('blockdevices'): if device['type'] == 'disk': disks_size.update({device['name']: device['size']}) return disks_size