import io
import logging
import mmap
from dataclasses import dataclass
from functools import singledispatchmethod
from typing import Dict, Optional, Type
from uuid import UUID
import archinfo
from cle.errors import CLEError, CLEUnknownFormatError
from . import Backend, register_backend
from .pe import PE
from .te import TE
try:
import uefi_firmware
except ImportError:
uefi_firmware = None
log = logging.getLogger(__name__)
[docs]class UefiDriverLoadError(Exception):
"""
This error is raised (and caught internally) if the data contained in the UEFI entity tree doesn't make sense.
"""
[docs]class UefiFirmware(Backend):
"""
A UEFI firmware blob loader. Support is provided by the ``uefi_firmware`` package.
"""
is_default = True
@classmethod
def _to_bytes(cls, fileobj: io.IOBase):
try:
fileno = fileobj.fileno()
except io.UnsupportedOperation:
pass
else:
return mmap.mmap(fileno, 0, access=mmap.ACCESS_READ)
if isinstance(fileobj, io.BytesIO):
return fileobj.getbuffer()
# fuck it, we'll do it live
fileobj.seek(0)
return fileobj.read()
[docs] @classmethod
def is_compatible(cls, stream):
if uefi_firmware is None:
return False
buffer = cls._to_bytes(stream)
parser = uefi_firmware.AutoParser(buffer)
return parser.type() != "unknown"
[docs] def __init__(self, *args, **kwargs) -> None:
if uefi_firmware is None:
raise CLEError("run `pip install uefi_firmware==1.10` to load UEFI firmware")
super().__init__(*args, **kwargs)
# hack: we are using a loader internal method in a non-kosher way which will cause our children to be
# marked as the main binary if we are also the main binary
# work around this by setting ourself here:
if self.loader._main_object is None:
self.loader._main_object = self
self._drivers: Dict[UUID, "UefiModuleMixin"] = {}
self._drivers_pending: Dict[UUID, "UefiModulePending"] = {}
self._current_file: Optional[UUID] = None
self.set_arch(archinfo.arch_from_id("x86_64")) # TODO: ???
buffer = self._to_bytes(self._binary_stream)
parser = uefi_firmware.AutoParser(buffer)
firmware = parser.parse()
self._load(firmware)
while self._drivers_pending:
uuid, pending = self._drivers_pending.popitem()
try:
child = pending.build(self, uuid)
except UefiDriverLoadError as e:
log.warning("Failed to load %s: %s", uuid, e.args[0])
else:
self._drivers[uuid] = child
child.parent_object = self
self.child_objects.append(child)
if self.child_objects:
self._arch = self.child_objects[0].arch
else:
log.warning("Loaded empty UEFI firmware?")
self.has_memory = False
self.pic = True
# hack pt. 2
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}")
if uefi_firmware is not None:
@_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)
[docs]@dataclass
class UefiModulePending:
"""
A worklist entry for the UEFI firmware loader.
"""
name: Optional[str] = None
pe_image: Optional[bytes] = None
te_image: Optional[bytes] = None
# version
# dependencies
[docs] def build(self, parent: UefiFirmware, guid: UUID) -> "UefiModuleMixin":
count = (self.pe_image is not None) + (self.te_image is not None)
if count > 1:
raise UefiDriverLoadError("Multiple image sections")
cls: "Type[UefiModuleMixin]"
if self.pe_image is not None:
cls = UefiPE
data = self.pe_image
elif self.te_image is not None:
cls = UefiTE
data = self.te_image
else:
raise UefiDriverLoadError("Missing PE or TE image section")
return cls(None, io.BytesIO(data), is_main_bin=False, loader=parent.loader, name=self.name, guid=guid)
[docs]class UefiModuleMixin(Backend):
"""
A mixin to make other kinds of backends load as UEFI modules.
"""
[docs] def __init__(self, *args, guid: UUID, name: Optional[str], **kwargs):
super().__init__(*args, **kwargs)
self.guid = guid
self.user_interface_name = name
if self.linked_base == 0:
self.pic = True
def __repr__(self):
return (
f"<{type(self).__name__} Object "
f'{self.guid}{f" {self.user_interface_name}" if self.user_interface_name else ""}, '
f"maps [{self.min_addr:#x}:{self.max_addr:#x}]>"
)
[docs]class UefiPE(UefiModuleMixin, PE):
"""
A PE file contained in a UEFI image.
"""
[docs]class UefiTE(UefiModuleMixin, TE):
"""
A TE file contained in a UEFI image.
"""
register_backend("uefi", UefiFirmware)