Setting EFI module entry function for module type

TE's use PEI module entry point type and PE's use DXE module entry
point type

Author: zznop

__init__.py

@@ -1,4 +1,5 @@
 from binaryninja import PluginCommand, BinaryView, BackgroundTaskThread, log_alert
+from .moduletype import identify_efi_module_type, set_efi_module_entry_type, EFIModuleType
 from .protocols import (
     init_protocol_mapping,
     define_handle_protocol_types,
@@ -20,34 +21,26 @@ def __init__(self, bv: BinaryView):
             super().__init__("Initializing EFI protocol mappings...", True)
             self.bv = bv
 
-        def run(self):
-            if not init_protocol_mapping():
+        def _resolve_dxe(self):
+            self.progress = "Propagating EFI system table pointers..."
+            if not propagate_system_table_pointers(self.bv, self):
                 return
 
-            if "EFI_SYSTEM_TABLE" not in self.bv.types:
-                log_alert("This binary is not using the EFI platform. Use Open with Options when loading the binary to select the EFI platform.")
+            self.progress = "Defining types for uses of HandleProtocol..."
+            if not define_handle_protocol_types(self.bv, self):
                 return
 
-            self.bv.begin_undo_actions()
-            try:
-                self.progress = "Propagating EFI system table pointers..."
-                if not propagate_system_table_pointers(self.bv, self):
-                    return
-
-                self.progress = "Defining types for uses of HandleProtocol..."
-                if not define_handle_protocol_types(self.bv, self):
-                    return
-
-                self.progress = "Defining types for uses of OpenProtocol..."
-                if not define_open_protocol_types(self.bv, self):
-                    return
+            self.progress = "Defining types for uses of OpenProtocol..."
+            if not define_open_protocol_types(self.bv, self):
+                return
 
-                self.progress = "Defining types for uses of LocateProtocol..."
-                if not define_locate_protocol_types(self.bv, self):
-                    return
+            self.progress = "Defining types for uses of LocateProtocol..."
+            if not define_locate_protocol_types(self.bv, self):
+                return
 
-                # SMM/MM types cannot be propagated until EFI_BOOT_SERVICES types are propagated and the
-                # EFI_MM_BASE_PROTOCOL or EFI_SMM_BASE2_PROTOCOL is resolved
+            # SMM/MM types cannot be propagated until EFI_BOOT_SERVICES types are propagated and the
+            # EFI_MM_BASE_PROTOCOL or EFI_SMM_BASE2_PROTOCOL is resolved
+            if "EFI_MM_SYSTEM_TABLE" in self.bv.types:
                 self.progress = "Defining types for SMM/MM system tables..."
                 if not define_locate_mm_system_table_types(self.bv, self) or not define_locate_smm_system_table_types(
                     self.bv, self
@@ -65,6 +58,32 @@ def run(self):
                     self.bv, self
                 ):
                     return
+
+        def run(self):
+            if not init_protocol_mapping():
+                return
+
+            if "EFI_SYSTEM_TABLE" not in self.bv.types:
+                log_alert("This binary is not using the EFI platform. Use Open with Options when loading the binary to select the EFI platform.")
+                return
+
+            module_type = identify_efi_module_type(self.bv)
+            if module_type == EFIModuleType.UNKNOWN:
+                log_alert("Could not identify the EFI module type.")
+                return
+
+            self.bv.begin_undo_actions()
+            try:
+                try:
+                    set_efi_module_entry_type(self.bv, module_type)
+                except SyntaxError:
+                    log_alert("Failed to set entry function type. Update Binary Ninja for latest EFI type definitions.")
+                    return
+
+                if module_type == EFIModuleType.DXE:
+                    self._resolve_dxe()
+                elif module_type == EFIModuleType.PEI:
+                    log_alert("TE PEI modules are not yet supported.")
             finally:
                 self.bv.commit_undo_actions()
 

moduletype.py

@@ -0,0 +1,46 @@
+from enum import Enum
+from binaryninja import BinaryView, SymbolType
+
+
+class EFIModuleType(Enum):
+    """UEFI module types"""
+
+    UNKNOWN = 0
+    DXE = 1
+    PEI = 2
+
+
+def set_efi_module_entry_type(bv: BinaryView, modtype: EFIModuleType) -> None:
+    """Set the prototype for the module entrypoint"""
+
+    _start = bv.get_symbol_by_raw_name("_start")
+    if not _start or _start.type != SymbolType.FunctionSymbol:
+        return
+
+    func = bv.get_function_at(_start.address)
+    if not func:
+        return
+
+    if modtype == EFIModuleType.PEI:
+        func.type = "EFI_STATUS _start(EFI_PEI_FILE_HANDLE FileHandle, EFI_PEI_SERVICES **PeiServices)"
+
+    if modtype == EFIModuleType.DXE:
+        func.type = "EFI_STATUS _start(EFI_HANDLE ImageHandle, EFI_SYSTEM_TABLE* SystemTable)"
+
+    func.reanalyze()
+
+
+def identify_efi_module_type(bv: BinaryView) -> EFIModuleType:
+    """Identify the type of EFI module
+
+    PE's are reported as DXE modules and TE's are reported as PEI modules. This is correct for most
+    cases, and is the convention used by most UEFI tooling.
+    """
+
+    if bv.get_view_of_type("PE"):
+        return EFIModuleType.DXE
+
+    if bv.get_view_of_type("TE"):
+        return EFIModuleType.PEI
+
+    return EFIModuleType.UNKNOWN