IDAPython script to apply enum values from MACRO_EFI to analysed EFI modules

efi-enums.py

# IDAPython script to apply enum values from MACRO_EFI to analysed EFI modules

from typing import Optional, List

import ida_allins
import ida_bytes
import ida_ida
import ida_typeinf
import ida_ua
import idaapi
import idautils


class EfiEnumsResolverError(Exception):
    """Generic enums resolver error."""

    def __init__(self, value: str) -> None:
        self.value = value

    def __str__(self):
        return repr(self.value)


class EfiEnumsResolver:
    TYPE_NAME = "MACRO_EFI"

    def __init__(self) -> None:
        self._load_uefi_til()
        self.mask = 0xFFFFFFFFFFFFFF00 if self._is_64_bits() else 0xFFFFFF00
        self.masked_value = 0x8000000000000000 if self._is_64_bits() else 0x80000000

        self._insn_types: Optional[List[int]] = None
        self._macro_efi_tid: Optional[int] = None

    def _load_uefi_til(self) -> None:
        ida_typeinf.add_til(
            "uefi64" if self._is_64_bits() else "uefi", ida_typeinf.ADDTIL_DEFAULT
        )

    @staticmethod
    def _is_64_bits() -> bool:
        return idaapi.idainfo_is_64bit()

    @staticmethod
    def _get_insn_types() -> List[int]:
        procname = ida_ida.inf_get_procname()
        if procname == "metapc":
            return [ida_allins.NN_mov]
        elif procname == "ARM":
            return [ida_allins.ARM_mov, ida_allins.ARM_movl]
        else:
            raise EfiEnumsResolverError(f"unsupported processor: {procname}")

    @property
    def insn_types(self) -> List[int]:
        if self._insn_types is None:
            self._insn_types = self._get_insn_types()
        return self._insn_types

    def _get_macro_efi_tid(self) -> int:
        idati = ida_typeinf.get_idati()
        tinfo = ida_typeinf.tinfo_t()
        if not tinfo.get_named_type(idati, EfiEnumsResolver.TYPE_NAME):
            raise EfiEnumsResolverError(
                f"{EfiEnumsResolver.TYPE_NAME} type is not found"
            )
        return tinfo.force_tid()

    @property
    def macro_efi_tid(self) -> int:
        if self._macro_efi_tid is None:
            self._macro_efi_tid = self._get_macro_efi_tid()
        return self._macro_efi_tid

    def resolve(self) -> bool:
        # enumerate all instructions, this can be optimised,
        # e.g. usually a value from MACRO_EFI is used in a
        # small basic blocks
        insn = idaapi.insn_t()
        for ea in idautils.Heads():
            idaapi.decode_insn(insn, ea)
            if insn.itype not in self.insn_types or insn.ops[0].type != ida_ua.o_reg:
                continue

            if insn.ops[1].value & self.mask != self.masked_value:
                continue

            print(f"fount enum value {insn.ops[1].value & self.mask:#x} at: {ea:#x}")
            ida_bytes.op_enum(ea, 1, self.macro_efi_tid, 0)

        return True


if __name__ == "__main__":
    try:
        resolver = EfiEnumsResolver()
        resolver.resolve()
    except EfiEnumsResolverError as e:
        print(f"error: {repr(e)}")