Add Structure Support for Ghidra Interface

yodalib/data/artifacts/struct.py

@@ -31,7 +31,7 @@ def __str__(self):
         return f"<StructMember: {self.type} {self.name}; @{hex(self.offset)}>"
 
     def __repr__(self):
-        self.__str__()
+        return self.__str__()
 
     @classmethod
     def parse(cls, s):

yodalib/decompilers/ghidra/interface.py

@@ -219,6 +219,49 @@ def _set_function_header(self, fheader: FunctionHeader, decompilation=None, **kw
 
         return changes
 
+    @ghidra_transaction
+    def _set_struct(self, struct: Struct, header=True, members=True, **kwargs) -> bool:
+        struct: Struct = struct
+        old_ghidra_struct = self._get_struct_by_name('/' + struct.name)
+        data_manager = self.ghidra.currentProgram.getDataTypeManager()
+        handler = self.ghidra.import_module_object("ghidra.program.model.data", "DataTypeConflictHandler")
+        structType = self.ghidra.import_module_object("ghidra.program.model.data", "StructureDataType")
+        byteType = self.ghidra.import_module_object("ghidra.program.model.data", "ByteDataType")
+        ghidra_struct = structType(struct.name, 0)
+        for offset in struct.members:
+            member = struct.members[offset]
+            ghidra_struct.add(byteType.dataType, 1, member.name, "")
+            ghidra_struct.growStructure(member.size - 1)
+            for dtc in ghidra_struct.getComponents():
+                if dtc.getFieldName() == member.name:
+                    gtype = self.typestr_to_gtype(member.type if member.type else 'undefined' + str(member.size))
+                    for i in range(offset, offset + member.size):
+                        ghidra_struct.clearAtOffset(i)
+                    ghidra_struct.replaceAtOffset(offset, gtype, member.size, member.name, "")
+                    break
+        try:
+            if old_ghidra_struct:
+                data_manager.replaceDataType(old_ghidra_struct, ghidra_struct, True)
+            else:
+                data_manager.addDataType(ghidra_struct, handler.DEFAULT_HANDLER)
+            return True
+        except Exception as ex:
+            print(f'Error filling struct {struct.name}: {ex}')
+            return False
+
+    def _get_struct(self, name) -> Optional[Struct]:
+        ghidra_struct = self._get_struct_by_name(name)
+        bs_struct = Struct(ghidra_struct.getName(), ghidra_struct.getLength(), self._struct_members_from_gstruct(name))
+        return bs_struct
+
+    def _structs(self) -> Dict[str, Struct]:
+        name_sizes: Optional[List[Tuple[str, int]]] = self.ghidra.bridge.remote_eval(
+            "[(s.getPathName(), s.getLength())"
+            "for s in currentProgram.getDataTypeManager().getAllStructures()]"
+        )
+        return {
+            name: Struct(name, size, members=self._struct_members_from_gstruct(name)) for name, size in name_sizes
+        } if name_sizes else {}
 
     #
     # TODO: REMOVE ME THIS IS THE BINSYNC CODE
@@ -291,34 +334,6 @@ def fill_comment(self, addr, user=None, artifact=None, **kwargs):
     # Artifact API
     #
 
-    def struct(self, name) -> Optional[Struct]:
-        ghidra_struct = self._get_struct_by_name(name)
-        members: Optional[List[Tuple[str, int, str, int]]] = self.ghidra.bridge.remote_eval(
-            "[(m.getFieldName(), m.getOffset(), m.getDataType().getName(), m.getLength()) if m.getFieldName() else "
-            "('field_'+hex(m.getOffset())[2:], m.getOffset(), m.getDataType().getName(), m.getLength()) "
-            "for m in ghidra_struct.getComponents()]",
-            ghidra_struct=ghidra_struct
-        )
-        struct_members = {}
-        if members:
-            struct_members = {
-                offset: StructMember(name, offset, typestr, size) for name, offset, typestr, size in members
-            }
-        bs_struct = Struct(ghidra_struct.getName(), ghidra_struct.getLength(), struct_members)
-        return bs_struct
-
-    def structs(self) -> Dict[str, Struct]:
-        name_sizes: Optional[List[Tuple[str, int]]] = self.ghidra.bridge.remote_eval(
-            "[(s.getPathName(), s.getLength())"
-            "for s in currentProgram.getDataTypeManager().getAllStructures()]"
-        )
-        structures = {}
-        if name_sizes:
-            structures = {
-                name: Struct(name, size, None) for name, size in name_sizes
-            }
-        return structures
-
     def global_var(self, addr) -> Optional[GlobalVariable]:
         light_global_vars = self.global_vars()
         for offset, global_var in light_global_vars.items():
@@ -392,6 +407,18 @@ def _get_local_variable_symbols(self, func: Function) -> Dict[str, "HighSymbol"]
     def _get_struct_by_name(self, name: str) -> "GhidraStructure":
         return self.ghidra.currentProgram.getDataTypeManager().getDataType(name)
 
+    def _struct_members_from_gstruct(self, name: str) -> Dict[int, StructMember]:
+        ghidra_struct = self._get_struct_by_name(name)
+        members: Optional[List[Tuple[str, int, str, int]]] = self.ghidra.bridge.remote_eval(
+            "[(m.getFieldName(), m.getOffset(), m.getDataType().getName(), m.getLength()) if m.getFieldName() else "
+            "('field_'+hex(m.getOffset())[2:], m.getOffset(), m.getDataType().getName(), m.getLength()) "
+            "for m in ghidra_struct.getComponents()]",
+            ghidra_struct=ghidra_struct
+        )
+        return {
+            offset: StructMember(name, offset, typestr, size) for name, offset, typestr, size in members
+        } if members else {}
+
     def _get_nearest_function(self, addr: int) -> "GhidraFunction":
         func_manager = self.ghidra.currentProgram.getFunctionManager()
         return func_manager.getFunctionContaining(self.ghidra.toAddr(addr))