Guard optional UEFI dispatch registration (#703)

* Guard optional UEFI dispatch registration

* Use merged PyVEX source

* Simplify optional UEFI backend loading

Author: zardus

cle/backends/uefi_firmware.py

@@ -4,7 +4,7 @@
 import logging
 import mmap
 from dataclasses import dataclass
-from functools import singledispatchmethod
+from typing import cast
 from uuid import UUID
 
 import archinfo
@@ -104,40 +104,36 @@ def __init__(self, *args, **kwargs) -> None:
         if self.loader._main_object is self:
             self.loader._main_object = None
 
-    @singledispatchmethod
-    def _load(self, uefi_obj):  # pylint: disable=no-self-use
-        raise CLEUnknownFormatError(f"Can't load firmware object: {uefi_obj}")
+    def _load(self, uefi_obj):
+        if uefi_obj is None:
+            return
+        if uefi_firmware is None:
+            raise ImportError("The UEFI backend requires the uefi-firmware package")
+
+        is_firmware_file = isinstance(uefi_obj, uefi_firmware.uefi.FirmwareFile)
+        old_uuid = self._current_file
+        if is_firmware_file:
+            if uefi_obj.type == 7:  # driver
+                uuid = UUID(bytes=uefi_obj.guid)
+                self._drivers_pending[uuid] = UefiModulePending()
+                self._current_file = uuid
+        elif isinstance(uefi_obj, uefi_firmware.uefi.FirmwareFileSystemSection):
+            pending = self._drivers_pending.get(self._current_file) if self._current_file is not None else None
+            if pending is not None:
+                if uefi_obj.type == 16:  # pe32 image
+                    pending.pe_image = cast(bytes, uefi_obj.content)
+                elif uefi_obj.type == 18:  # te image
+                    pending.te_image = cast(bytes, uefi_obj.content)
+                elif uefi_obj.type == 21:  # user interface name
+                    pending.name = cast(bytes, uefi_obj.content).decode("utf-16").strip("\0")
+        elif not isinstance(uefi_obj, uefi_firmware.FirmwareObject):
+            raise CLEUnknownFormatError(f"Can't load firmware object: {uefi_obj}")
 
-    @_load.register
-    def _load_generic(self, uefi_obj: uefi_firmware.FirmwareObject):
         for obj in uefi_obj.objects:
             self._load(obj)
 
-    @_load.register
-    def _load_none(self, uefi_obj: None):
-        pass
-
-    @_load.register
-    def _load_firmwarefile(self, uefi_obj: uefi_firmware.uefi.FirmwareFile):
-        old_uuid = self._current_file
-        if uefi_obj.type == 7:  # driver
-            uuid = UUID(bytes=uefi_obj.guid)
-            self._drivers_pending[uuid] = UefiModulePending()
-            self._current_file = uuid
-        self._load_generic(uefi_obj)
-        self._current_file = old_uuid
-
-    @_load.register
-    def _load_firmwarefilesection(self, uefi_obj: uefi_firmware.uefi.FirmwareFileSystemSection):
-        pending = self._drivers_pending.get(self._current_file, None)
-        if pending is not None:
-            if uefi_obj.type == 16:  # pe32 image
-                pending.pe_image = uefi_obj.content
-            elif uefi_obj.type == 18:  # te image
-                pending.te_image = uefi_obj.content
-            elif uefi_obj.type == 21:  # user interface name
-                pending.name = uefi_obj.content.decode("utf-16").strip("\0")
-        self._load_generic(uefi_obj)
+        if is_firmware_file:
+            self._current_file = old_uuid
 
 
 @dataclass
@@ -184,7 +180,7 @@ def __init__(self, *args, guid: UUID, name: str | None, **kwargs):
     def __repr__(self):
         return (
             f"<{type(self).__name__} Object "
-            f'{self.guid}{f" {self.user_interface_name}" if self.user_interface_name else ""}, '
+            f"{self.guid}{f' {self.user_interface_name}' if self.user_interface_name else ''}, "
             f"maps [{self.min_addr:#x}:{self.max_addr:#x}]>"
         )
 

tests/test_optional_backends.py

@@ -0,0 +1,18 @@
+from __future__ import annotations
+
+import subprocess
+import sys
+
+
+def test_import_without_uefi_firmware():
+    script = """
+import sys
+
+sys.modules["uefi_firmware"] = None
+
+import cle
+from cle.backends import ALL_BACKENDS
+
+assert "uefi" in ALL_BACKENDS
+"""
+    subprocess.run([sys.executable, "-c", script], check=True)