# # Copyright (C) 2024 Soleta Networks # # 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)