account-for-letter-case-in-python-code-completion

Ghidra/Features/Base/src/main/java/ghidra/app/plugin/core/console/CodeCompletion.java

@@ -24,17 +24,28 @@ import javax.swing.JComponent;
  * 
  * It is intended to be used by the code completion process, especially the
  * CodeCompletionWindow.  It encapsulates:
- * - a description of the completion (what are you completing?)
+ * <ul>
- * - the actual String that will be inserted
+ * <li> a description of the completion (what are you completing?)
- * - an optional Component that will be in the completion List
+ * <li> the actual String that will be inserted
- * 
+ * <li> an optional Component that will be in the completion List
- * 
+ * <li> the number of characters to remove before the insertion of the completion
- *
+ * </ul>
+ * <p>
+ * For example, if one wants to autocomplete a string "Runscr" into "runScript", 
+ * the fields may look as follows:
+ * <ul>
+ * <li> description: "runScript (Method)"
+ * <li> insertion: "runScript"
+ * <li> component: null or JLabel("runScript (Method)")
+ * <li> charsToRemove: 6 (i.e. the length of "Runscr", 
+ *      as it may be required later to correctly replace the string)
+ * </ul>
  */
 public class CodeCompletion implements Comparable<CodeCompletion> {
 	private String description;
 	private String insertion;
 	private JComponent component;
+	private int charsToRemove;
 
 	
 	/**
@@ -60,6 +71,24 @@ public class CodeCompletion implements Comparable<CodeCompletion> {
 		this.description = description;
 		this.insertion = insertion;
 		this.component = comp;
+		this.charsToRemove = 0;
+	}
+
+
+	/**
+	 * Construct a new CodeCompletion.
+	 * 
+	 * @param description description of this completion
+	 * @param insertion what will be inserted (or null)
+	 * @param comp (optional) Component to appear in completion List (or null)
+	 * @param charsToRemove the number of characters that should be removed before the insertion
+	 */
+	public CodeCompletion(String description, String insertion,
+			JComponent comp, int charsToRemove) {
+		this.description = description;
+		this.insertion = insertion;
+		this.component = comp;
+		this.charsToRemove = charsToRemove;
 	}
 
 	
@@ -94,6 +123,16 @@ public class CodeCompletion implements Comparable<CodeCompletion> {
 		return insertion;
 	}
 	
+	/**
+	 * Returns the number of characters to remove from the input before the insertion
+	 * of the code completion
+	 * 
+	 * @return the number of characters to remove
+	 */
+	public int getCharsToRemove() {
+		return charsToRemove;
+	}
+	
 	
 	/**
 	 * Returns a String representation of this CodeCompletion.

Ghidra/Features/Base/src/main/java/ghidra/app/plugin/core/interpreter/InterpreterPanel.java

@@ -624,16 +624,21 @@ public class InterpreterPanel extends JPanel implements OptionsChangeListener {
 		String insertion = completion.getInsertion();
 
 		/* insert completion string */
-		setInputTextPaneText(text.substring(0, position) + insertion + text.substring(position));
+		int insertedTextStart = position - completion.getCharsToRemove();
+		int insertedTextEnd = insertedTextStart + insertion.length();
+		var inputText = text.substring(0, insertedTextStart) + insertion + text.substring(position);
+		setInputTextPaneText(inputText);
 
 		/* Select what we inserted so that the user can easily
 		 * get rid of what they did (in case of a mistake). */
 		if (highlightCompletion) {
-			inputTextPane.setSelectionStart(position);
+			inputTextPane.setSelectionStart(insertedTextStart);
+			inputTextPane.moveCaretPosition(insertedTextEnd);
+		} else {
+			/* Then put the caret right after what we inserted. */
+			inputTextPane.setCaretPosition(insertedTextEnd);
 		}
-
+	
-		/* Then put the caret right after what we inserted. */
-		inputTextPane.moveCaretPosition(position + insertion.length());
 		updateCompletionList();
 	}
 

Ghidra/Features/Base/src/test.slow/java/ghidra/app/plugin/core/interpreter/InterpreterPanelTest.java

@@ -17,6 +17,7 @@ package ghidra.app.plugin.core.interpreter;
 
 import static org.junit.Assert.*;
 
+import java.awt.event.KeyEvent;
 import java.io.*;
 import java.util.List;
 import java.util.concurrent.CountDownLatch;
@@ -48,6 +49,7 @@ public class InterpreterPanelTest extends AbstractGhidraHeadedIntegrationTest {
 	private JTextPane inputTextPane;
 	private Document inputDoc;
 	private BufferedReader reader;
+	private List<CodeCompletion> testingCodeCompletions = List.of();	
 
 	@Before
 	public void setUp() throws Exception {
@@ -152,6 +154,37 @@ public class InterpreterPanelTest extends AbstractGhidraHeadedIntegrationTest {
 
 		assertTrue(ip.getStdin().available() > 0);
 	}
+	
+	@Test(timeout = 20000)
+	public void testCodeCompletionInsertion() {
+		// test the general ability to insert text with the code completion popup;
+		// and also make sure that the completion doesn't accidentally destroy 
+		// a part of the input near the caret
+		
+		// case 1: "simple." => "simple.completion"
+		testingCodeCompletions = List.of(
+				new CodeCompletion("test", "completion", null, 0));
+		triggerText(inputTextPane, "simple.");
+		insertFirstCodeCompletion();
+		assertEquals("simple.completion", inputTextPane.getText());
+		triggerEnter(inputTextPane);
+		
+		// case 2: "simple." => "not.so.simple.completion"
+		testingCodeCompletions = List.of(
+				new CodeCompletion("test", "not.so.simple.completion", null, "simple.".length()));
+		triggerText(inputTextPane, "simple.");
+		insertFirstCodeCompletion();
+		assertEquals("not.so.simple.completion", inputTextPane.getText());
+		triggerEnter(inputTextPane);
+		
+		// case 3: "( check.both.sides.of<caret> )" => "( check.is.ok )"
+		testingCodeCompletions = List.of(
+				new CodeCompletion("test", "is.ok", null, "both.sides.of".length()));
+		triggerText(inputTextPane, "( check.both.sides.of )");
+		inputTextPane.setCaretPosition("( check.both.sides.of".length());
+		insertFirstCodeCompletion();
+		assertEquals("( check.is.ok )", inputTextPane.getText());
+	}
 
 	private InterpreterPanel createIP() {
 		InterpreterConnection dummyIC = new InterpreterConnection() {
@@ -167,7 +200,7 @@ public class InterpreterPanelTest extends AbstractGhidraHeadedIntegrationTest {
 
 			@Override
 			public List<CodeCompletion> getCompletions(String cmd) {
-				return List.of();
+				return testingCodeCompletions;
 			}
 		};
 
@@ -271,4 +304,11 @@ public class InterpreterPanelTest extends AbstractGhidraHeadedIntegrationTest {
 			}
 		}).start();
 	}
+	
+	private void insertFirstCodeCompletion() {
+		KeyStroke defaultCompletionTrigger = CompletionWindowTrigger.TAB.getKeyStroke();
+		triggerKey(inputTextPane, defaultCompletionTrigger);
+		triggerActionKey(inputTextPane, 0, KeyEvent.VK_DOWN);
+		triggerEnter(inputTextPane);
+	}
 }

Ghidra/Features/Python/python-src/introspect.py

@@ -74,8 +74,9 @@ def getAutoCompleteList(command='', locals=None, includeMagic=1,
                     pyObj = locals[attribute]
                 completion_list.append(PythonCodeCompletionFactory.
                                        newCodeCompletion(attribute,
-                                                         attribute[len(filter):],
+                                                         attribute, 
-                                                         pyObj))
+                                                         pyObj,
+                                                         filter))
             except:
                 # hmm, problem evaluating?  Examples of this include
                 # inner classes, e.g. access$0, which aren't valid Python

Ghidra/Features/Python/src/main/java/ghidra/python/PythonCodeCompletionFactory.java

@@ -183,9 +183,28 @@ public class PythonCodeCompletionFactory {
 	 * @param insertion what will be inserted to make the code complete
 	 * @param pyObj a Python Object
 	 * @return A new CodeCompletion from the given Python objects.
+	 * @deprecated use {@link #newCodeCompletion(String, String, PyObject, String)} instead,
+	 *             it allows creation of substituting code completions
 	 */
+	@Deprecated
 	public static CodeCompletion newCodeCompletion(String description, String insertion,
 			PyObject pyObj) {
+		return newCodeCompletion(description, insertion, pyObj, "");
+	}
+	
+	/**
+	 * Creates a new CodeCompletion from the given Python objects.
+	 * 
+	 * @param description description of the new CodeCompletion
+	 * @param insertion what will be inserted to make the code complete
+	 * @param pyObj a Python Object
+	 * @param userInput a word we want to complete, can be an empty string.
+	 *        It's used to determine which part (if any) of the input should be 
+	 *        removed before the insertion of the completion
+	 * @return A new CodeCompletion from the given Python objects.
+	 */
+	public static CodeCompletion newCodeCompletion(String description, String insertion,
+			PyObject pyObj, String userInput) {
 		JComponent comp = null;
 
 		if (pyObj != null) {
@@ -213,7 +232,9 @@ public class PythonCodeCompletionFactory {
 				}
 			}
 		}
-		return new CodeCompletion(description, insertion, comp);
+		
+		int charsToRemove = userInput.length();
+		return new CodeCompletion(description, insertion, comp, charsToRemove);
 	}
 
 	/**

Ghidra/Features/Python/src/test.slow/java/ghidra/python/PythonCodeCompletionTest.java

@@ -0,0 +1,147 @@
+/* ###
+ * IP: GHIDRA
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ * 
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ * 
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package ghidra.python;
+
+import static org.junit.Assert.*;
+
+import java.io.File;
+import java.io.IOException;
+import java.nio.charset.Charset;
+import java.util.Collection;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+import org.apache.commons.io.FileUtils;
+import org.junit.*;
+import org.junit.rules.TemporaryFolder;
+
+import generic.jar.ResourceFile;
+import ghidra.app.plugin.core.console.CodeCompletion;
+import ghidra.app.plugin.core.osgi.BundleHost;
+import ghidra.app.script.GhidraScriptUtil;
+import ghidra.test.AbstractGhidraHeadedIntegrationTest;
+
+/**
+ * Tests for the Ghidra Python Interpreter's code completion functionality.
+ */
+public class PythonCodeCompletionTest extends AbstractGhidraHeadedIntegrationTest {
+
+	private String simpleTestProgram = """
+			my_int = 32
+			my_bool = True
+			my_string = 'this is a string'
+			my_list = ["a", 2, 5.3, my_string]
+			my_tuple = (1, 2, 3)
+			my_dictionary = {"key1": "1", "key2": 2, "key3": my_list}
+			mY_None = None
+			i = 5
+
+			def factorial(n):
+			    return 1 if n == 0 else n * factorial(n-1)
+			def error_function():
+			    raise IOError("An IO error occurred!")
+
+			class Employee:
+			    def __init__(self, id, name):
+			        self.id = id
+			        self.name = name
+			    def getId(self):
+			        return self.id
+			    def getName(self):
+			        return self.name
+
+			employee = Employee(42, "Bob")
+				""".stripIndent();
+
+	@Rule
+	public TemporaryFolder tempScriptFolder = new TemporaryFolder();
+
+	private GhidraPythonInterpreter interpreter;
+
+	@Before
+	public void setUp() throws Exception {
+		GhidraScriptUtil.initialize(new BundleHost(), null);
+		interpreter = GhidraPythonInterpreter.get();
+		executePythonProgram(simpleTestProgram);
+	}
+
+	@After
+	public void tearDown() throws Exception {
+		interpreter.cleanup();
+		GhidraScriptUtil.dispose();
+	}
+
+	@Test
+	public void testBasicCodeCompletion() {
+		// test the "insertion" field
+		// it should be equal to the full name of a variable we want to complete
+
+		List<String> completions = List.of("my_bool", "my_dictionary", "my_int",
+				"my_list", "mY_None", "my_string", "my_tuple");
+		assertCompletionsInclude("My", completions);
+		assertCompletionsInclude("employee.Get", List.of("getId", "getName"));
+		assertCompletionsInclude("('noise', (1 + fact", List.of("factorial"));
+	}
+
+	@Test
+	public void testCharsToRemoveField() {
+		// 'charsToRemove' field should be equal to the length of
+		// a part of variable/function/method name we are trying to complete here.
+		// This allows us to correctly put a completion in cases when we really
+		// just want to replace a piece of text (i.e. "CURRENTAddress" => "currentAddress")
+		// rather than simply 'complete' it.
+
+		assertCharsToRemoveEqualsTo("my_int", "my_int".length());
+		assertCharsToRemoveEqualsTo("employee.get", "get".length());
+		assertCharsToRemoveEqualsTo("('noise', (1 + fact", "fact".length());
+
+		assertCharsToRemoveEqualsTo("employee.", 0);
+		assertCharsToRemoveEqualsTo("employee.getId(", 0);
+	}
+
+	private void assertCompletionsInclude(String command, 
+			Collection<String> expectedCompletions) {
+		Set<String> completions = interpreter.getCommandCompletions(command, false)
+				.stream()
+				.map(c -> c.getInsertion())
+				.collect(Collectors.toSet());
+
+		var missing = new HashSet<String>(expectedCompletions);
+		missing.removeAll(completions);
+		if (!missing.isEmpty()) {
+			Assert.fail("Could't find these completions: " + missing);
+		}
+	}
+
+	private void assertCharsToRemoveEqualsTo(String command, int expectedCharsToRemove) {
+		for (CodeCompletion comp : interpreter.getCommandCompletions(command, false)) {
+			assertEquals(String.format("%s; field 'charsToRemove' ", comp),
+					expectedCharsToRemove, comp.getCharsToRemove());
+		}
+	}
+
+	private void executePythonProgram(String code) {
+		try {
+			File tempFile = tempScriptFolder.newFile();
+			FileUtils.writeStringToFile(tempFile, code, Charset.defaultCharset());
+			interpreter.execFile(new ResourceFile(tempFile), null);
+		} catch (IOException e) {
+			fail("couldn't create a test script: " + e.getMessage());
+		}
+	}
+}