enum progress

binsync · Nov 21, 2023 · f3a6a7a · f3a6a7a
1 parent 00aba3b
commit f3a6a7a
Showing 1 changed file with 20 additions and 2 deletions.
diff --git a/yodalib/decompilers/ghidra/interface.py b/yodalib/decompilers/ghidra/interface.py
@@ -270,9 +270,27 @@ 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:
-        # TODO: Implement me
-        return False
+        old_ghidra_enum = self.ghidra.currentProgram.getDataTypeManager().getDataType(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")
+        path = categoryPath(enum.name)
+        outer_path = categoryPath(path.getPath())
+        ghidra_enum = enumType(outer_path, path.getName(), 4)
+        for member_name in enum.members.keys():
+            ghidra_enum.add(member_name, enum.members[member_name])
+        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]:
         return Enum(name, self._get_enum_members(name))