Support classes, enums and add/remove

bridge_mcp_ghidra.py

@@ -435,19 +435,183 @@ def search_bytes(bytes_hex: str, offset: int = 0, limit: int = 100) -> list:
     )
 
 @mcp.tool()
-def set_bytes(address: str, bytes_hex: str) -> str:
+def create_enum(name: str, category: str = None, size: int = 4, values: list = None) -> str:
     """
-    Writes a sequence of bytes to the specified address in the program's memory.
+    Create a new enum.
+
+    Args:
+        name: The name of the new enum.
+        category: The category path for the enum (e.g., /my_enums). Defaults to root.
+        size: The size of the enum in bytes (default: 4).
+        values: A list of value dictionaries to add to the new enum.
+                Each dict should have 'name', 'value', and optionally 'comment'.
+                Example: [{"name": "VALUE1", "value": 0, "comment": "First value"}]
+
+    Returns:
+        A status message indicating success or failure.
+    """
+    data = {"name": name, "size": str(size)}
+    if category:
+        data["category"] = category
+    if values:
+        data["values"] = json.dumps(values)
+    return safe_post("create_enum", data)
+
+@mcp.tool()
+def add_enum_values(enum_name: str, values: list, category: str = None) -> str:
+    """
+    Add values to an existing enum.
+
+    Args:
+        enum_name: The name of the enum to modify.
+        values: A list of value dictionaries to add to the enum.
+                Each dict should have 'name', 'value', and optionally 'comment'.
+                Example: [{"name": "VALUE1", "value": 0, "comment": "First value"}]
+        category: The category path for the enum. Defaults to root.
+
+    Returns:
+        A status message indicating success or failure.
+    """
+    data = {"enum_name": enum_name, "values": json.dumps(values)}
+    if category:
+        data["category"] = category
+    return safe_post("add_enum_values", data)
+
+@mcp.tool()
+def get_enum(name: str, category: str = None) -> dict:
+    """
+    Get an enum's definition.
+
+    Args:
+        name: The name of the enum.
+        category: The category path for the enum. Defaults to root.
+
+    Returns:
+        A dictionary representing the enum, or an error message.
+    """
+    params = {"name": name}
+    if category:
+        params["category"] = category
+
+    response_lines = safe_get("get_enum", params)
+    response_str = "\n".join(response_lines)
+
+    try:
+        # Attempt to parse the JSON response
+        return json.loads(response_str)
+    except json.JSONDecodeError:
+        # If it's not JSON, it's likely an error message
+        return {"error": response_str}
+
+@mcp.tool()
+def set_global_data_type(address: str, data_type: str, length: int = -1, clear_mode: str = "CHECK_FOR_SPACE") -> str:
+    """
+    Set the data type of a global variable or data at a specific memory address.
+
+    Args:
+        address: The memory address in hex format (e.g., "0x401000")
+        data_type: The name of the data type to apply (e.g., "int", "char*", "MyStruct")
+        length: Optional length for dynamic data types (default: -1, let type determine)
+        clear_mode: How to handle conflicting data. Options:
+                   - "CHECK_FOR_SPACE": Ensure data fits before clearing (default)
+                   - "CLEAR_SINGLE_DATA": Always clear single code unit at address
+                   - "CLEAR_ALL_UNDEFINED_CONFLICT_DATA": Clear conflicting undefined data
+                   - "CLEAR_ALL_DEFAULT_CONFLICT_DATA": Clear conflicting default data
+                   - "CLEAR_ALL_CONFLICT_DATA": Clear all conflicting data
+
+    Returns:
+        A status message indicating success or failure.
+    """
+    data = {
+        "address": address,
+        "data_type": data_type,
+        "clear_mode": clear_mode
+    }
+    if length > 0:
+        data["length"] = str(length)
+
+    return safe_post("set_global_data_type", data)
 
+@mcp.tool()
+def add_class_members(class_name: str, members: list, parent_namespace: str = None) -> str:
+    """
+    Add members to an existing C++ class.
+
     Args:
-        address: Destination address (e.g., "0x140001000")
-        bytes_hex: Sequence of space-separated bytes in hexadecimal format (e.g., "90 90 90 90")
+        class_name: The name of the class to modify.
+        members: A list of member dictionaries to add to the class.
+                Each dict should have 'name', 'type', and optionally 'offset' and 'comment'.
+                Example: [{"name": "health", "type": "float", "comment": "Player health"}]
+        parent_namespace: The parent namespace where the class is located (optional).
+
+    Returns:
+        A status message indicating success or failure.
+    """
+    params = {"class_name": class_name, "members": json.dumps(members)}
+    if parent_namespace:
+        params["parent_namespace"] = parent_namespace
+
+    return safe_post("add_class_members", params)
 
+@mcp.tool()
+def remove_class_members(class_name: str, members: list, parent_namespace: str = None) -> str:
+    """
+    Remove members from an existing C++ class.
+
+    Args:
+        class_name: The name of the class to modify.
+        members: A list of member names to remove from the class.
+                Example: ["old_member", "unused_field"]
+        parent_namespace: The parent namespace where the class is located (optional).
+
     Returns:
-        Result of the operation (e.g., "Bytes written successfully" or a detailed error)
+        A status message indicating success or failure.
     """
-    return safe_post("set_bytes", {"address": address, "bytes": bytes_hex})
+    params = {"class_name": class_name, "members": json.dumps(members)}
+    if parent_namespace:
+        params["parent_namespace"] = parent_namespace
+
+    return safe_post("remove_class_members", params)
+
+@mcp.tool()
+def remove_enum_values(enum_name: str, values: list, category: str = None) -> str:
+    """
+    Remove values from an existing enum.
+
+    Args:
+        enum_name: The name of the enum to modify.
+        values: A list of value names to remove from the enum.
+                Example: ["OLD_VALUE", "DEPRECATED_OPTION"]
+        category: The category path for the enum (optional, defaults to root).
+
+    Returns:
+        A status message indicating success or failure.
+    """
+    params = {"enum_name": enum_name, "values": json.dumps(values)}
+    if category:
+        params["category"] = category
+
+    return safe_post("remove_enum_values", params)
+
+@mcp.tool()
+def remove_struct_members(struct_name: str, members: list, category: str = None) -> str:
+    """
+    Remove members from an existing struct.
+
+    Args:
+        struct_name: The name of the struct to modify.
+        members: A list of member names to remove from the struct.
+                Example: ["old_field", "unused_member"]
+        category: The category path for the struct (optional, defaults to root).
+
+    Returns:
+        A status message indicating success or failure.
+    """
+    params = {"struct_name": struct_name, "members": json.dumps(members)}
+    if category:
+        params["category"] = category
 
+    return safe_post("remove_struct_members", params)
 
 def main():
     parser = argparse.ArgumentParser(description="MCP server for Ghidra")

src/main/java/com/lauriewired/handlers/act/AddClassMembers.java

@@ -0,0 +1,185 @@
+package com.lauriewired.handlers.act;
+
+import com.lauriewired.handlers.Handler;
+import com.sun.net.httpserver.HttpExchange;
+import ghidra.framework.plugintool.PluginTool;
+import ghidra.program.model.data.DataType;
+import ghidra.program.model.data.DataTypeComponent;
+import ghidra.program.model.data.DataTypeManager;
+import ghidra.program.model.data.Structure;
+import ghidra.program.model.listing.GhidraClass;
+import ghidra.program.model.listing.Program;
+import ghidra.program.model.symbol.Namespace;
+import ghidra.program.model.symbol.Symbol;
+import ghidra.program.model.symbol.SymbolTable;
+import ghidra.program.model.symbol.SymbolType;
+import ghidra.program.database.data.DataTypeUtilities;
+
+import com.google.gson.Gson;
+
+import javax.swing.SwingUtilities;
+import java.io.IOException;
+import java.lang.reflect.InvocationTargetException;
+import java.util.*;
+import java.util.concurrent.atomic.AtomicReference;
+
+import static com.lauriewired.util.GhidraUtils.*;
+import static com.lauriewired.util.ParseUtils.*;
+import static com.lauriewired.util.StructUtils.StructMember;
+import static ghidra.program.util.GhidraProgramUtilities.getCurrentProgram;
+
+/**
+ * Handler for adding members to a C++ class in Ghidra.
+ * This modifies both the class namespace and its associated structure data type.
+ * Expects a POST request with parameters:
+ * - class_name: Name of the class to modify
+ * - parent_namespace: Parent namespace where the class is located (optional)
+ * - members: JSON array of members to add, each with fields:
+ *   - type: Data type of the member
+ *   - name: Name of the member
+ *   - comment: Comment for the member (optional)
+ *   - offset: Offset in bytes (optional, -1 for next available position)
+ */
+public final class AddClassMembers extends Handler {
+	/**
+	 * Constructor for the AddClassMembers handler.
+	 *
+	 * @param tool The Ghidra plugin tool instance.
+	 */
+	public AddClassMembers(PluginTool tool) {
+		super(tool, "/add_class_members");
+	}
+
+	/**
+	 * Handles the HTTP request to add members to a class.
+	 *
+	 * @param exchange The HTTP exchange containing the request and response.
+	 * @throws IOException If an I/O error occurs during handling.
+	 */
+	@Override
+	public void handle(HttpExchange exchange) throws IOException {
+		Map<String, String> params = parsePostParams(exchange);
+		String className = params.get("class_name");
+		String parentNamespace = params.get("parent_namespace");
+		String membersJson = params.get("members");
+
+		if (className == null || membersJson == null) {
+			sendResponse(exchange, "class_name and members are required");
+			return;
+		}
+		sendResponse(exchange, addClassMembers(className, parentNamespace, membersJson));
+	}
+
+	/**
+	 * Adds members to a class in the current Ghidra program.
+	 *
+	 * @param className The name of the class to modify.
+	 * @param parentNamespace The parent namespace where the class is located (optional).
+	 * @param membersJson JSON array of members to add.
+	 * @return A message indicating success or failure.
+	 */
+	private String addClassMembers(String className, String parentNamespace, String membersJson) {
+		Program program = getCurrentProgram(tool);
+		if (program == null)
+			return "No program loaded";
+
+		final AtomicReference<String> result = new AtomicReference<>();
+		try {
+			SwingUtilities.invokeAndWait(() -> {
+				int txId = program.startTransaction("Add Class Members");
+				boolean success = false;
+				try {
+					SymbolTable symbolTable = program.getSymbolTable();
+					DataTypeManager dtm = program.getDataTypeManager();
+
+					// Find the class namespace
+					Namespace parent = program.getGlobalNamespace();
+					if (parentNamespace != null && !parentNamespace.isEmpty()) {
+						parent = symbolTable.getNamespace(parentNamespace, program.getGlobalNamespace());
+						if (parent == null) {
+							result.set("Error: Parent namespace '" + parentNamespace + "' not found");
+							return;
+						}
+					}
+
+					// Find the class by iterating through symbols
+					GhidraClass classNamespace = null;
+					for (Symbol symbol : symbolTable.getSymbols(className, parent)) {
+						if (symbol.getSymbolType() == SymbolType.CLASS) {
+							classNamespace = (GhidraClass) symbol.getObject();
+							break;
+						}
+					}
+
+					if (classNamespace == null) {
+						result.set("Error: Class '" + className + "' not found" + 
+								(parent != null ? " in namespace " + parent.getName() : ""));
+						return;
+					}
+
+					// Find the associated structure
+					Structure classStruct = DataTypeUtilities.findExistingClassStruct(dtm, classNamespace);
+					if (classStruct == null) {
+						result.set("Error: No structure found for class '" + className + "'");
+						return;
+					}
+
+					StringBuilder responseBuilder = new StringBuilder(
+							"Adding members to class " + className);
+
+					if (membersJson != null && !membersJson.isEmpty()) {
+						Gson gson = new Gson();
+						StructMember[] members = gson.fromJson(membersJson, StructMember[].class);
+
+						int membersAdded = 0;
+						for (StructMember member : members) {
+							DataType memberDt = resolveDataType(tool, dtm, member.type);
+							if (memberDt == null) {
+								responseBuilder.append("\nError: Could not resolve data type '").append(member.type)
+										.append("' for member '").append(member.name)
+										.append("'. Aborting further member creation.");
+								break;
+							}
+
+							// Check if member already exists
+							DataTypeComponent existingComponent = null;
+							for (DataTypeComponent comp : classStruct.getComponents()) {
+								if (comp.getFieldName() != null && comp.getFieldName().equals(member.name)) {
+									existingComponent = comp;
+									break;
+								}
+							}
+							if (existingComponent != null) {
+								responseBuilder.append("\nWarning: Member '").append(member.name)
+										.append("' already exists. Skipping.");
+								continue;
+							}
+
+							if (member.offset != -1) {
+								classStruct.insertAtOffset((int) member.offset, memberDt, -1, member.name, member.comment);
+								responseBuilder.append("\nAdded member '").append(member.name)
+										.append("' at offset ").append(member.offset);
+							} else {
+								classStruct.add(memberDt, member.name, member.comment);
+								responseBuilder.append("\nAdded member '").append(member.name)
+										.append("' at end of structure");
+							}
+							membersAdded++;
+						}
+						responseBuilder.append("\nSuccessfully added ").append(membersAdded).append(" members to class ").append(className);
+						result.set(responseBuilder.toString());
+						success = membersAdded > 0;
+					}
+
+				} catch (Exception e) {
+					result.set("Error: Failed to add members to class: " + e.getMessage());
+				} finally {
+					program.endTransaction(txId, success);
+				}
+			});
+		} catch (InterruptedException | InvocationTargetException e) {
+			return "Error: Failed to execute add class members on Swing thread: " + e.getMessage();
+		}
+		return result.get();
+	}
+}