Add Enum Support to Ghidra Decompiler Interface (#3)

* Template * Enums finished * Fixed _artifact_getter definition for enums * enum progress * FINALLY got enum setting working * Deleted old binsync code * Fix non-generic enum setting and some comment stuff * Sped up enum getting and finished name fixes * Added error handling to get_enum --------- Co-authored-by: mahaloz <[email protected]>
binsync · Nov 29, 2023 · f681372 · f681372
1 parent 7540641
commit f681372
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 3 deletions.
diff --git a/yodalib/api/artifact_dict.py b/yodalib/api/artifact_dict.py
@@ -21,7 +21,7 @@ def __init__(self, artifact_cls, decompiler_interface: "DecompilerInterface", er
             StackVariable: (self._di._set_stack_variable, self._di._get_stack_variable, self._di._stack_variables),
             GlobalVariable: (self._di._set_global_variable, self._di._get_global_var, self._di._global_vars),
             Struct: (self._di._set_struct, self._di._get_struct, self._di._structs),
-            Enum: (self._di._set_enum, self._di._enums, self._di._enums),
+            Enum: (self._di._set_enum, self._di._get_enum, self._di._enums),
             Comment: (self._di._set_comment, self._di._get_comment, self._di._comments),
             Patch: (self._di._set_patch, self._di._get_patch, self._di._patches)
         }

diff --git a/yodalib/decompilers/ghidra/interface.py b/yodalib/decompilers/ghidra/interface.py
@@ -6,7 +6,7 @@
 from yodalib.api import DecompilerInterface
 from yodalib.api.decompiler_interface import requires_decompilation
 from yodalib.data import (
-    Function, FunctionHeader, StackVariable, Comment, FunctionArgument, GlobalVariable, Struct, StructMember
+    Function, FunctionHeader, StackVariable, Comment, FunctionArgument, GlobalVariable, Struct, StructMember, Enum
 )
 
 from .artifact_lifter import GhidraArtifactLifter
@@ -270,6 +270,40 @@ def _structs(self) -> Dict[str, Struct]:
             name: Struct(name, size, members=self._struct_members_from_gstruct(name)) for name, size in name_sizes
         } if name_sizes else {}
 
+    @ghidra_transaction
+    def _set_enum(self, enum: Enum, **kwargs) -> bool:
+        corrected_enum_name = "/" + enum.name
+        old_ghidra_enum = self.ghidra.currentProgram.getDataTypeManager().getDataType(corrected_enum_name)
+        data_manager = self.ghidra.currentProgram.getDataTypeManager()
+        handler = self.ghidra.import_module_object("ghidra.program.model.data", "DataTypeConflictHandler")
+        enumType = self.ghidra.import_module_object("ghidra.program.model.data", "EnumDataType")
+        categoryPath = self.ghidra.import_module_object("ghidra.program.model.data", "CategoryPath")
+        ghidra_enum = enumType(categoryPath('/'), enum.name, 4)
+        for m_name, m_val in enum.members.items():
+            ghidra_enum.add(m_name, m_val)
+
+        try:
+            if old_ghidra_enum:
+                data_manager.replaceDataType(old_ghidra_enum, ghidra_enum, True)
+            else:
+                data_manager.addDataType(ghidra_enum, handler.DEFAULT_HANDLER)
+            return True
+        except Exception as ex:
+            print(f'Error adding enum {enum.name}: {ex}')
+            return False
+
+    def _get_enum(self, name) -> Optional[Enum]:
+        members = self._get_enum_members('/' + name)
+        return Enum(name, members) if members else None
+
+    def _enums(self) -> Dict[str, Enum]:
+        names: Optional[List[str]] = self.ghidra.bridge.remote_eval(
+            "[dType.getPathName() "
+            "for dType in currentProgram.getDataTypeManager().getAllDataTypes()"
+            "if str(type(dType)) == \"<type 'ghidra.program.database.data.EnumDB'>\"]"
+        )
+        return {name[1:]: Enum(name[1:], self._get_enum_members(name)) for name in names if name.count('/') == 1} if names else {}
+
     #
     # TODO: REMOVE ME THIS IS THE BINSYNC CODE
     # Filler/Setter API
@@ -327,6 +361,8 @@ def fill_comment(self, addr, user=None, artifact=None, **kwargs):
         code_unit = self.ghidra.import_module_object("ghidra.program.model.listing", "CodeUnit")
         set_cmt_cmd_cls = self.ghidra.import_module_object("ghidra.app.cmd.comments", "SetCommentCmd")
         cmt_type = code_unit.PRE_COMMENT if comment.decompiled else code_unit.EOL_COMMENT
+        if comment.addr == comment.func_addr:
+            cmt_type = code_unit.PLATE_COMMENT
 
         if comment.comment:
             # TODO: check if comment already exists, and append?
@@ -340,7 +376,6 @@ def fill_comment(self, addr, user=None, artifact=None, **kwargs):
     # TODO: REMOVE ME THIS IS ALSO BINSYNC CODE
     # Artifact API
     #
-
     def global_var(self, addr) -> Optional[GlobalVariable]:
         light_global_vars = self.global_vars()
         for offset, global_var in light_global_vars.items():
@@ -460,6 +495,17 @@ def _struct_members_from_gstruct(self, name: str) -> Dict[int, StructMember]:
             offset: StructMember(name, offset, typestr, size) for name, offset, typestr, size in members
         } if members else {}
 
+    def _get_enum_members(self, name: str) -> Optional[Dict[str, int]]:
+        ghidra_enum = self.ghidra.currentProgram.getDataTypeManager().getDataType(name)
+        if not ghidra_enum:
+            return None
+        name_vals: Optional[List[Tuple[str, int]]] = self.ghidra.bridge.remote_eval(
+            "[(name, ghidra_enum.getValue(name))"
+            "for name in ghidra_enum.getNames()]",
+            ghidra_enum=ghidra_enum
+        )
+        return {name: value for name, value in name_vals} if name_vals else {}
+
     def _get_nearest_function(self, addr: int) -> "GhidraFunction":
         func_manager = self.ghidra.currentProgram.getFunctionManager()
         return func_manager.getFunctionContaining(self.ghidra.toAddr(addr))