grub: replace legacy grub install scripts

Translate old legacy grub scripts into grub.py
Implement ogGrubInstallMbr as install_main_grub() and
ogGrubInstallPartition as install_linux_grub().

Add grub configuration file generator through the classes
GrubConfig and MenuEntry.

Ensure EFI tree structure compatibility with legacy code.
The structure of the created folders in the ESP is non-standard,
efi binaries are usually located in the folder below the EFI/
directory.

Structure used by ogClient:

EFI/
├── grub/
│   └── Boot/
│       ├── BOOTX64.CSV
│       ├── grub.cfg
│       ├── mmx64.efi
│       ├── shimx64.efi
│       ├── BOOTX64.EFI
│       ├── grubx64.efi
│       └── ogloader.efi
...

The function _mangle_efi_folder handles the folder structure after
grub-install to comply with the location expected by ogLive.

install_linux_grub() installs a grub local to each Linux install
to enable chainloading, each grub is located in EFI/Part-xx-yy/ in
UEFI. The local linux BIOS grub in legacy scripts is unreliable,
grub-install reports a failure during the install process.

install_main_grub() installs a global grub in EFI/grub/ to show a
grub menu when the pxe boot fails. The global grub contains entries
to every installed os. No global grub is installed for BIOS
systems, a Boot partition would be required to store the grub
configuration.

4 files changed, 401 insertions, 31 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)
diff --git a/src/utils/postinstall.py b/src/utils/postinstall.py
index 221da81..90b8cc6 100644
--- a/src/utils/postinstall.py
+++ b/src/utils/postinstall.py
@@ -12,6 +12,7 @@ import logging
 import hivex
 import shlex
 from src.log import OgError
+from src.utils.grub import install_main_grub, install_linux_grub
 from src.utils.bcd import update_bcd
 from src.utils.probe import *
 from src.utils.disk import *
@@ -152,19 +153,6 @@ def configure_mbr_boot_sector(disk, partition):
         logging.warning(f'{cmd_configure} returned non-zero exit status {proc.returncode}')
 
 
-def configure_grub_in_mbr(disk, partition):
-    cmd_configure = f"ogGrubInstallMbr {disk} {partition} TRUE"
-
-    proc = subprocess.run(cmd_configure,
-                          stdout=subprocess.DEVNULL,
-                          stderr=subprocess.DEVNULL,
-                          encoding='utf-8',
-                          shell=True,
-                          check=True)
-    if proc.returncode != 0:
-        logging.warning(f'{cmd_configure} returned non-zero exit status {proc.returncode}')
-
-
 def configure_fstab(disk, partition):
     logging.info(f'Configuring /etc/fstab')
     device = get_partition_device(disk, partition)
@@ -178,31 +166,17 @@ def configure_fstab(disk, partition):
     finally:
         umount(mountpoint)
 
-
-def install_grub(disk, partition):
-    cmd_configure = f"ogGrubInstallPartition {disk} {partition}"
-
-    proc = subprocess.run(cmd_configure,
-                          stdout=subprocess.DEVNULL,
-                          stderr=subprocess.DEVNULL,
-                          encoding='utf-8',
-                          shell=True,
-                          check=True)
-    if proc.returncode != 0:
-        logging.warning(f'{cmd_configure} returned non-zero exit status {proc.returncode}')
-
-
 def configure_os_linux(disk, partition):
     hostname = gethostname()
     set_linux_hostname(disk, partition, hostname)
 
     configure_fstab(disk, partition)
 
+    install_linux_grub(disk, partition)
+
     if is_uefi_supported():
         _, _, esp_part_number = get_efi_partition(disk, enforce_gpt=True)
-        configure_grub_in_mbr(disk, esp_part_number)
-
-    install_grub(disk, partition)
+        install_main_grub()
 
 
 def configure_os_windows(disk, partition):
@@ -213,7 +187,7 @@ def configure_os_windows(disk, partition):
         restore_windows_efi_bootloader(disk, partition)
 
         _, _, esp_part_number = get_efi_partition(disk, enforce_gpt=True)
-        configure_grub_in_mbr(disk, esp_part_number)
+        install_main_grub()
     else:
         configure_mbr_boot_sector(disk, partition)
 
diff --git a/src/utils/probe.py b/src/utils/probe.py
index e53f08f..92ccf7c 100644
--- a/src/utils/probe.py
+++ b/src/utils/probe.py
@@ -10,6 +10,7 @@ import os
 import subprocess
 import platform
 import logging
+import sys
 
 from enum import Enum
 from subprocess import PIPE
@@ -76,6 +77,10 @@ def getwindowsversion(winreghives):
     return 'Microsoft Windows'
 
 
+def interpreter_is64bit():
+    return sys.maxsize > 2**32
+
+
 def windows_is64bit(winreghives):
     """
     Check for 64 bit Windows by means of retrieving the value of
diff --git a/src/utils/uefi.py b/src/utils/uefi.py
index 27b7ba0..10095a7 100644
--- a/src/utils/uefi.py
+++ b/src/utils/uefi.py
@@ -125,6 +125,24 @@ def efibootmgr_create_bootentry(disk, part, loader, label, add_to_bootorder=True
     except OSError as e:
         raise OgError(f'Unexpected error adding boot entry to nvram. UEFI firmware might be buggy') from e
 
+def efibootmgr_set_entry_order(label, position):
+    logging.info(f'Setting {label} entry to position {position} of boot order')
+    boot_info = run_efibootmgr_json(validate=False)
+    boot_entries = boot_info.get('vars', [])
+    boot_order = boot_info.get('BootOrder', [])
+
+    entry = _find_bootentry(boot_entries, label)
+    target_grub_entry = _strip_boot_prefix(entry)
+
+    if target_grub_entry in boot_order:
+        boot_order.remove(target_grub_entry)
+
+    boot_order.insert(position, target_grub_entry)
+
+    try:
+        proc = subprocess.run([EFIBOOTMGR_BIN, "-o", ",".join(boot_order)], check=True, text=True)
+    except OSError as e:
+        raise OgError(f'Unexpected error setting boot order to NVRAM. UEFI firmware might be buggy') from e
 
 def _find_efi_loader(loader_paths):
     for efi_app in loader_paths: