diff options
Diffstat (limited to 'src/utils/grub.py')
| -rw-r--r-- | src/utils/grub.py | 373 |
1 files changed, 373 insertions, 0 deletions
diff --git a/src/utils/grub.py b/src/utils/grub.py new file mode 100644 index 0000000..ca7eb86 --- /dev/null +++ b/src/utils/grub.py @@ -0,0 +1,373 @@ +# +# Copyright (C) 2024 Soleta Networks <info@soleta.eu> +# +# This program is free software: you can redistribute it and/or modify it under +# the terms of the GNU Affero General Public License as published by the +# Free Software Foundation; either version 3 of the License, or +# (at your option) any later version. + +import subprocess +import logging +import shutil +import shlex +import os + +from enum import Enum +from src.utils.probe import * +from src.utils.bios import * +from src.utils.uefi import * +from src.utils.disk import * +from src.log import OgError + +GRUB_TIMEOUT = 5 + +class GrubConfig: + def __init__(self, timeout): + self.timeout = timeout + self.entries = [] + + def add_entry(self, entry): + self.entries.append(entry) + + def __str__(self): + res = f'set timeout={self.timeout}\n' + res += f'set default=0\n' + res += '\n' + for entry in self.entries: + res += str(entry) + res += '\n' + return res + +class MenuEntry: + class PartType(Enum): + MBR = 'part_msdos' + GPT = 'part_gpt' + + class FsType(Enum): + EXT = 'ext2' + FAT = 'fat' + + def __init__(self, name, disk_num, part_num, part_type, fs): + self.name = self._escape_menuentry_name(name) + self.disk_num = disk_num + self.part_num = part_num + self.part_type = part_type + self.fs = fs + self.efipath = None + self.initrd = None + self.vmlinuz = None + self.vmlinuz_root = None + self.params = None + + def set_chainload_data(self, efipath): + self.efipath = efipath + + def set_linux_data(self, initrd, vmlinuz, vmlinuz_root, params): + self.initrd = initrd + self.vmlinuz = vmlinuz + self.vmlinuz_root = vmlinuz_root + self.params = params + + def _set_root_str(self): + if self.part_type == MenuEntry.PartType.GPT: + return f'set root=(hd{self.disk_num - 1},gpt{self.part_num})' + else: + return f'set root=(hd{self.disk_num - 1},{self.part_num})' + + def _get_chainload_str(self): + res = f'menuentry "{self.name}" {{\n' + res += f' insmod {self.part_type.value}\n' + res += ' insmod chain\n' + res += f' insmod {self.fs.value}\n' + res += f' {self._set_root_str()}\n' + res += f' chainloader {self.efipath}\n' + res += '}\n' + return res + + def _get_linux_str(self): + res = f'menuentry "{self.name}" {{\n' + res += f' insmod {self.part_type.value}\n' + res += ' insmod linux\n' + res += f' insmod {self.fs.value}\n' + res += f' {self._set_root_str()}\n' + res += f' linux {self.vmlinuz} root={self.vmlinuz_root} {self.params}\n' + res += f' initrd {self.initrd}\n' + res += '}\n' + return res + + def __str__(self): + if self.efipath: + return self._get_chainload_str() + else: + return self._get_linux_str() + + @staticmethod + def _escape_menuentry_name(entry_name): + entry_name = entry_name.replace('"', r'\"') + entry_name = entry_name.replace('\\', r'\\') + entry_name = entry_name.replace('$', r'\$') + entry_name = entry_name.replace('*', r'\*') + entry_name = entry_name.replace('!', r'\!') + return entry_name + +def _get_linux_data(disk_num, part_num, mountpoint): + os_entry = {} + os_entry['name'] = f'{os_probe(mountpoint)} ({disk_num}, {part_num})' + os_entry['device'] = get_partition_device(disk_num, part_num) + os_entry['part_type'] = get_disk_part_type(disk_num) + + kernel_path = get_vmlinuz_path(mountpoint) + os_entry['vmlinuz'] = '/' + os.path.relpath(kernel_path, mountpoint) + initrd_path = get_initrd_path(mountpoint) + os_entry['initrd'] = '/' + os.path.relpath(initrd_path, mountpoint) + return os_entry + +def _generate_linux_grub_config(disk_num, part_num, mountpoint, grub_cfg): + grub_config = GrubConfig(timeout=0) + + os_entry = _get_linux_data(disk_num, part_num, mountpoint) + part_type = get_disk_part_type(disk_num) + + menu_entry = MenuEntry(name=os_entry['name'], + disk_num=disk_num, + part_num=part_num, + part_type=part_type, + fs=MenuEntry.FsType.EXT) + menu_entry.set_linux_data(initrd=os_entry['initrd'], + vmlinuz=os_entry['vmlinuz'], + vmlinuz_root=os_entry['device'], + params='ro quiet splash') + grub_config.add_entry(menu_entry) + + with open(grub_cfg, 'w') as f: + f.write(str(grub_config)) + +def _mangle_efi_folder(entry_dir, boot_dir): + efi_boot_src = f'{entry_dir}/EFI/BOOT' + for file_name in os.listdir(efi_boot_src): + shutil.move(f'{efi_boot_src}/{file_name}', boot_dir) + + shutil.rmtree(f'{entry_dir}/EFI') + + shutil.copyfile("/usr/lib/shim/shimx64.efi.signed", + f'{boot_dir}/shimx64.efi') + + shutil.copyfile(f'{boot_dir}/grubx64.efi', + f'{boot_dir}/ogloader.efi') + +def _install_linux_grub_efi(disk_num, part_num, device, mountpoint): + if interpreter_is64bit(): + arch = 'x86_64-efi' + else: + logging.warning("Old 32-bit UEFI system found here") + arch = 'i386-efi' + _esp_disk_num = 1 + + esp, _esp_disk, _esp_part_number = get_efi_partition(_esp_disk_num, enforce_gpt=False) + esp_mountpoint = esp.replace('dev', 'mnt') + + if not mount_mkdir(esp, esp_mountpoint): + raise OgError(f'Unable to mount detected EFI System Partition at {esp} into {esp_mountpoint}') + + _bootlabel = f'Part-{disk_num:02d}-{part_num:02d}' + entry_dir = f'{esp_mountpoint}/EFI/{_bootlabel}' + boot_dir = f'{entry_dir}/Boot' + if os.path.exists(entry_dir): + shutil.rmtree(entry_dir) + os.makedirs(boot_dir) + + logging.info(f'Calling grub-install with target {arch} in {entry_dir}') + grub_install_cmd = (f'grub-install --removable --no-nvram --uefi-secure-boot ' + f'--target={arch} --efi-directory={entry_dir} ' + f'--root-directory={entry_dir} --recheck') + + try: + subprocess.run(shlex.split(grub_install_cmd), check=True) + except subprocess.CalledProcessError as e: + umount(esp_mountpoint) + raise OgError(f"Error during GRUB install: {e}") from e + + _mangle_efi_folder(entry_dir, boot_dir) + + logging.info(f'Generating grub.cfg in {boot_dir}') + grub_cfg = f'{boot_dir}/grub.cfg' + try: + _generate_linux_grub_config(disk_num, part_num, mountpoint, grub_cfg) + except Exception as e: + raise OgError(f'Error generating {grub_cfg}: {e}') from e + finally: + umount(esp_mountpoint) + +def _install_linux_grub_bios(disk_num, part_num, device, mountpoint): + arch = 'i386-pc' + + entry_dir = f'{mountpoint}/boot' + grub_dir = f'{entry_dir}/grub' + if os.path.exists(grub_dir): + shutil.rmtree(grub_dir) + os.makedirs(grub_dir) + + logging.info(f'Calling grub-install with target {arch} in {entry_dir}') + grub_install_cmd = f'grub-install --force --target={arch} --boot-directory={entry_dir} --recheck {device}' + + try: + subprocess.run(shlex.split(grub_install_cmd), check=True) + except subprocess.CalledProcessError as e: + umount(mountpoint) + raise OgError(f"Error during GRUB install: {e}") from e + + grub_cfg = f'{grub_dir}/grub.cfg' + + logging.info(f'Generating grub.cfg in {grub_dir}') + try: + _generate_linux_grub_config(disk_num, part_num, mountpoint, grub_cfg) + except Exception as e: + raise OgError(f'Error generating {grub_cfg}: {e}') from e + finally: + umount(mountpoint) + +def install_linux_grub(disk_num, part_num): + device = get_partition_device(disk_num, part_num) + + mountpoint = device.replace('dev', 'mnt') + if not mount_mkdir(device, mountpoint): + raise OgError(f'Unable to mount {device} into {mountpoint}') + + try: + if is_uefi_supported(): + logging.info(f'Installing GRUB for UEFI Linux at {device}...') + _install_linux_grub_efi(disk_num, part_num, device, mountpoint) + else: + logging.info(f'Installing GRUB for BIOS Linux at {device}...') + _install_linux_grub_bios(disk_num, part_num, device, mountpoint) + finally: + umount(mountpoint) + +def _get_os_entries(esp_mountpoint): + os_entries = [] + available_disks = get_disks() + for disk_num, diskname in enumerate(available_disks, start=1): + disk_device = f'/dev/{diskname}' + partitions_data = get_partition_data(device=disk_device) + + for p in partitions_data: + part_num = p.partno + 1 + mountpoint = p.padev.replace('dev', 'mnt') + + if mountpoint == esp_mountpoint: + continue + + if not mount_mkdir(p.padev, mountpoint): + raise OgError(f'Unable to mount {p.padev} into {mountpoint}') + + try: + os_family = get_os_family(mountpoint) + system_name = os_probe(mountpoint) + except Exception as e: + umount(mountpoint) + raise + + if os_family == OSFamily.UNKNOWN: + umount(mountpoint) + continue + + os_entry = {} + os_entry['name'] = f'{system_name} ({disk_num}, {part_num})' + + _bootlabel = f'Part-{disk_num:02d}-{part_num:02d}' + try: + if os_family == OSFamily.WINDOWS: + efi_loader = find_windows_efi_loader(esp_mountpoint, _bootlabel) + elif os_family == OSFamily.LINUX: + efi_loader = find_linux_efi_loader(esp_mountpoint, _bootlabel) + os_entry['efipath'] = '/' + os.path.relpath(efi_loader, esp_mountpoint) + finally: + umount(mountpoint) + + os_entries.append(os_entry) + return os_entries + +def _generate_main_grub_config(grub_cfg, esp_disk_num, esp_part_num, esp_mountpoint): + os_entries = _get_os_entries(esp_mountpoint) + + esp_part_type = get_disk_part_type(disk_num=1) + + grub_config = GrubConfig(timeout=GRUB_TIMEOUT) + for os_entry in os_entries: + menu_entry = MenuEntry(name=os_entry['name'], + disk_num=esp_disk_num, + part_num=esp_part_num, + part_type=esp_part_type, + fs=MenuEntry.FsType.FAT) + menu_entry.set_chainload_data(efipath=os_entry['efipath']) + grub_config.add_entry(menu_entry) + + with open(grub_cfg, 'w') as f: + f.write(str(grub_config)) + +def get_disk_part_type(disk_num): + device = get_disks()[disk_num - 1] + cxt = fdisk.Context(f'/dev/{device}') + return MenuEntry.PartType.MBR if cxt.label.name == 'dos' else MenuEntry.PartType.GPT + +def _update_nvram(esp_disk, esp_part_number): + loader_path = '/EFI/grub/Boot/shimx64.efi' + bootlabel = 'grub' + efibootmgr_delete_bootentry(bootlabel) + efibootmgr_create_bootentry(esp_disk, esp_part_number, loader_path, bootlabel, add_to_bootorder=False) + efibootmgr_set_entry_order(bootlabel, 1) + +def install_main_grub(): + disk_device = f'/dev/{get_disks()[0]}' + is_uefi = is_uefi_supported() + if not is_uefi: + logging.info(f'Global GRUB install not supported in legacy BIOS') + return + + logging.info(f'Installing GRUB at {disk_device}') + _esp_disk_num = 1 + esp, _esp_disk, _esp_part_number = get_efi_partition(_esp_disk_num, enforce_gpt=False) + esp_mountpoint = esp.replace('dev', 'mnt') + + if not mount_mkdir(esp, esp_mountpoint): + raise OgError(f'Unable to mount detected EFI System Partition at {esp} into {esp_mountpoint}') + + entry_dir = f'{esp_mountpoint}/EFI/grub' + boot_dir = f'{entry_dir}/Boot' + if os.path.exists(entry_dir): + shutil.rmtree(entry_dir) + os.makedirs(boot_dir) + + if interpreter_is64bit(): + arch = 'x86_64-efi' + else: + logging.warning("Old 32-bit UEFI system found here") + arch = 'i386-efi' + logging.info(f'Calling grub-install with target {arch} in {entry_dir}') + grub_install_cmd = (f'grub-install --removable --no-nvram --uefi-secure-boot ' + f'--target={arch} --efi-directory={entry_dir} ' + f'--root-directory={entry_dir} --recheck') + + try: + subprocess.run(shlex.split(grub_install_cmd), check=True) + except subprocess.CalledProcessError as e: + umount(esp_mountpoint) + raise OgError(f"Error during GRUB install: {e}") from e + + _mangle_efi_folder(entry_dir, boot_dir) + + logging.info(f'Generating grub.cfg in {boot_dir}') + grub_cfg = f'{boot_dir}/grub.cfg' + try: + _generate_main_grub_config(grub_cfg, _esp_disk_num, _esp_part_number, esp_mountpoint) + except Exception as e: + umount(esp_mountpoint) + raise OgError(f'Error generating {grub_cfg}: {e}') from e + + logging.info('Updating grub UEFI NVRAM entry') + try: + _update_nvram(_esp_disk, _esp_part_number) + except Exception as e: + logging.info(f'Error updating NVRAM: {e}') + + umount(esp_mountpoint) |