From 8f37f93bf5c12bb8d4a794c51b9006f6e1da6df1 Mon Sep 17 00:00:00 2001 From: Florian Kothmeier Date: Sat, 9 Sep 2023 20:08:00 +0200 Subject: [PATCH] Add basic autocompletion --- data/python/complete.py | 268 +++++++++++++++++ data/python/jepwrappers.py | 37 +-- .../GhidrathonConsoleInputThread.java | 125 +++++++- .../java/ghidrathon/GhidrathonPlugin.java | 25 +- .../PythonCodeCompletionFactory.java | 281 ++++++++++++++++++ .../interpreter/GhidrathonInterpreter.java | 40 +++ 6 files changed, 741 insertions(+), 35 deletions(-) create mode 100644 data/python/complete.py create mode 100644 src/main/java/ghidrathon/PythonCodeCompletionFactory.java diff --git a/data/python/complete.py b/data/python/complete.py new file mode 100644 index 0000000..10dec85 --- /dev/null +++ b/data/python/complete.py @@ -0,0 +1,268 @@ +import ast +import inspect +import types + +import java + +from ghidrathon import PythonCodeCompletionFactory + + +def isJavaModule(module): + modules = ['ghidra.', 'java.'] + return any(module.startswith(name) for name in modules) + + +def isJavaMethod(obj): + "Returns whether the given object is a bound method implemented in Java" + return (isinstance(obj, types.MethodType) and hasattr(obj, '__self__') and hasattr(obj.__self__, '__module__') and isJavaModule(obj.__self__.__module__)) or obj.__class__.__name__ == "PythonCodeCompletionFactory$InspectableJavaMethod" + + +class CompletionObject: + "Object returned by getObject. See the documentation of getObject" + def __init__(self, o): + self.obj = o + + def getmembers(self): + "Returns the properties of the encapsulated object as a (name, value) tuple" + if self.obj.__class__.__name__ == "PythonCodeCompletionFactory$InspectableJavaObject": + return self.obj.getProperties() + return inspect.getmembers(self.obj) + + +def getObject(value, locals): + """ + Attempts to resolve the object that would be generated when evaluating + the AST `value` within the local variables `locals`. + This method does not run any python code and thus does not produce + side effects. + This also means we may not have enough information to determine which + object would be produced by this expression. + + This function returns a CompletionObject if the object produced by this + expression could be found. If not, this function returns None. + + Note: In case a method call on a Java object is processed, we return + a `PythonCodeCompletionFactory$InspectableJavaObject` that fakes the + properties and methods of the Java Object that would be returned if + called. This is necessary, as calling the function could produce side + effects. + + On a similar note, for python functions we also cannot return the actual + object that would be returned. In this case we return the returntype of + the function signature (if available). This should mostly work but may + be missing properties that are dynamically assigned on such an object + during creation. But this is currently the best we can do for python. + """ + match value: + case ast.Constant(literal): + return CompletionObject(literal) + case ast.Name(id, ctx): + try: + # Get object + obj = eval(id, locals) + return CompletionObject(obj) + except NameError: + return None + case ast.Call(func, _, _): + prop = getObject(func, locals) + retval = None + + # Hack to introspect java methods + if prop and isJavaMethod(prop.obj): + retval = PythonCodeCompletionFactory.getReturnType(prop.obj.__self__, prop.obj.__name__) + if retval and retval.getSrcClass() == java.lang.String: + # jep autoconverts between basic types + # There are more of these edge cases but this one happens the most + retval = '' + # If it ain't java, we may have a python signature + # if not, then we are lost + elif prop and inspect.signature(prop.obj).return_annotation: + retval = inspect.signature(prop.obj).return_annotation + if retval: + return CompletionObject(retval) + return None + case ast.Attribute(value, attr, _): + prop = getObject(value, locals) + props = [y for (x, y) in prop.getmembers() if x == attr] + if props: + # There may be multiple functions with different signatures + # This only happens when those members are reported through + # PythonCodeCompletionFactory$InspectableJavaObject.getProperties() + # and not when inspected via python + # This is bad for our cause, but we just have to live with that + # Just pick one, we don't check the signature anyways + # And luckily they all have the same return type + return CompletionObject(props[0]) + return None + case ast.Subscript | ast.ListComp | ast.SetComp | ast.GeneratorExp | ast.DictComp: + # TODO, can we handle this? + return None + case default: + raise ValueError(f"I don't know how to handle '{ast.dump(default)}' (getObject)") + + +def getProperties(value, locals): + """ + Returns a list of properties of the AST given by `value` when evaluated + within the local variables `locals` + Each entry of the returned list is of the form (name, prop). + `name` is a string and `prop` is the value of that property. + Because we do not actually run this code, we may not have enough + information to offer introspection for this value. In this case + the resulting list is empty + """ + prop = getObject(value, locals) + if prop: + return prop.getmembers() + return [] + + +def getVariables(locals): + """ + Returns all variables in the local scope and the builtins. + That is because most of the jepwrappers are bound to the builtins... + And thus don't show up in the local scope + """ + return [(x, locals[x]) for x in locals if locals[x] != locals] + \ + [(x, __builtins__[x]) for x in __builtins__] + + +def makeCompletions(values, prefix=''): + """ + Returns a list of CodeCompletion objects for all properties in the list + that start with the given prefix. + """ + return [PythonCodeCompletionFactory.newCodeCompletion(name, name[len(prefix):], value) + for name, value in values if name.startswith(prefix)] + + +def completeAST(parsed, locals, needs_property): + """ + Tries to provide autocompletion suggestions for the AST given by `parsed` + with the variables in scope given by locals. + Due to the way we have to handle property access, we need + a special case when the original input ended with a trailing point. + Therefore if `needs_property` is True, we return the properties of + the object returned if the given AST were to be evaluated. + If `needs_property` is False instead, we treat the last property access + in the AST as unfinished and report all properties that start with the last + property name as a prefix. + See the complete function for more information + """ + match parsed: + case ast.Constant(literal): + if needs_property: + return makeCompletions(getProperties(parsed, locals)) + return [] + case ast.Expr(value): + return completeAST(value, locals, needs_property) + case ast.UnaryOp(_, operand): + return completeAST(operand, locals, needs_property) + case ast.BinOp(_, _, right): + return completeAST(right, locals, needs_property) + case ast.BoolOp(_, values): + return completeAST(values[-1], locals, needs_property) + case ast.Compare(_, _, comparators): + return completeAST(comparators, locals, needs_property) + case ast.Name(id, _): + if needs_property: + return makeCompletions(getProperties(parsed, locals)) + return makeCompletions(getVariables(locals), id) + case ast.Call(func, _, _): + if needs_property: + return makeCompletions(getProperties(parsed, locals)) + # This is a valid function call and not an unfinished fragment + # There is nothing to complete + return [] + case ast.IfExp(_, _, orelse): + return completeAST(orelse, locals, needs_property) + case ast.Attribute(value, attr, _): + if needs_property: + # We need to complete a full property at the end + return makeCompletions(getProperties(parsed, locals), '') + # We already typed part of the property, let's see what we could complete + return makeCompletions(getProperties(value, locals), attr) + case ast.NamedExpr(): + # Of the form (x := 4). If this is parsed, then it is already complete + # Therefore there's nothing to complete here + return [] + case ast.Subscript() | ast.ListComp() | ast.SetComp() | ast.GeneratorExp() | ast.DictComp(): + if needs_property: + # Just pass introspection to getProperties + return makeCompletions(getProperties(parsed, locals), '') + # This is a valid expression and not an unfinished fragment + # There is nothing to complete + return [] + case ast.Starred(value, _): + return completeAST(value, locals, needs_property) + case ast.Assign(_, value, _): + return completeAST(value, locals, needs_property) + case ast.AnnAssign(_, _, value, _): + return completeAST(value, locals, needs_property) + case ast.AugAssign(_, _, value): + return completeAST(value, locals, needs_property) + case ast.Raise(exc, cause): + if cause: + return completeAST(cause, locals, needs_property) + return completeAST(exc, locals, needs_property) + case ast.Assert(test, msg): + if msg: + return completeAST(msg, locals, needs_property) + return completeAST(test, locals, needs_property) + case ast.Delete(targets): + return completeAST(targets[-1], locals, needs_property) + case ast.Pass | ast.Break | ast.Continue: + return [] + case ast.Return(value): + if value: + return completeAST(value, locals, needs_property) + return [] + case ast.Lambda(_, body): + return completeAST(body, locals, needs_property) + case ast.Yield(value) | ast.YieldFrom(value): + return completeAST(body, locals, needs_property) + case ast.Import() | ast.ImportFrom(): + # TODO, import autocompletion? + return [] + case ast.Global() | ast.Nonlocal(): + # TODO, autocomplete variable names? + return [] + # All other cases should be multi-line expressions, which our + # interpreter console does not support + case default: + raise ValueError(f"I don't know how to handle '{ast.dump(default)}' (completeAST)") + + +def complete(cmd, locals): + """ + Tries to provide autocompletion suggestions for the input string `cmd` + with the variables in scope given by locals. + """ + # Python with a trailing point will never compile + # We have to handle this special case seperately + needs_property = cmd.endswith('.') + # Special case: floating point literals, e.g. + # 0. is parsed as a float in python ;) + if len(cmd) > 2 and cmd[-2] in '0123456789': + needs_property = False + if needs_property: + cmd = cmd[:-1] + + # CMD may not be valid python, therefore we + # get the longest suffix that is syntactically correct + # This should work most of the time + parsed = None + while True: + # Uh oh, there is no valid expression + if not cmd: + # Just return a list of the locals + return makeCompletions(getVariables(locals)) + + try: + parsed = ast.parse(cmd, mode='single') + break + except SyntaxError: + cmd = cmd[1:].lstrip() + # Parsed will always be of type ast.Interactive + # We only have to complete the last expression + return completeAST(parsed.body[-1], locals, needs_property) diff --git a/data/python/jepwrappers.py b/data/python/jepwrappers.py index 1d12a44..14f334a 100644 --- a/data/python/jepwrappers.py +++ b/data/python/jepwrappers.py @@ -11,9 +11,10 @@ import os import java.lang +import ghidra +from ghidrathon import PythonCodeCompletionFactory cache_key = "ghidrathon_cache" -flatprogramapi_wrapper_stub = """@flatprogramapi_wrapper\ndef %s(*args, **kwargs): ...""" class GhidrathonCachedStream: @@ -98,10 +99,9 @@ def remove_state(): del get_cache()[get_java_thread_id()] -def flatprogramapi_wrapper(api): - def wrapped(*args, **kwargs): - return getattr(get_script(), api.__name__)(*args, **kwargs) - +def flatprogramapi_wrapper(attr, retval): + def wrapped(*args, **kwargs) -> retval: + return getattr(get_script(), attr)(*args, **kwargs) return wrapped @@ -293,41 +293,44 @@ def wrap_flatprogramapi_functions(): if not callable(attr_o): continue - # dynamically generate wrapper stub using attribute name - exec(flatprogramapi_wrapper_stub % attr, globals()) - - # add dynamically generated wrapper stub to __builtins__ - __builtins__[attr] = globals()[attr] + retval = PythonCodeCompletionFactory.getReturnTypeForClass(ghidra.app.script.GhidraScript, attr) + # Special case: Jep converts between basic types automatically + # There are more special cases like this, but this is the most frequent one + if retval == java.lang.String: + retval = str + # dynamically generate wrapper stub using attribute name and return value + # and add dynamically generated wrapper stub to __builtins__ + __builtins__[attr] = flatprogramapi_wrapper(attr, retval) wrap_flatprogramapi_functions() -def wrapped_monitor(): +def wrapped_monitor() -> ghidra.util.task.TaskMonitor: return get_script().getMonitor() -def wrapped_state(): +def wrapped_state() -> ghidra.app.script.GhidraState: return get_script_state() -def wrapped_currentProgram(): +def wrapped_currentProgram() -> ghidra.program.model.listing.Program: return get_script_state().getCurrentProgram() -def wrapped_currentAddress(): +def wrapped_currentAddress() -> ghidra.program.model.address.Address: return get_script_state().getCurrentAddress() -def wrapped_currentLocation(): +def wrapped_currentLocation() -> ghidra.program.util.ProgramLocation: return get_script_state().getCurrentLocation() -def wrapped_currentSelection(): +def wrapped_currentSelection() -> ghidra.program.util.ProgramSelection: return get_script_state().getCurrentSelection() -def wrapped_currentHighlight(): +def wrapped_currentHighlight() -> ghidra.program.util.ProgramSelection: return get_script_state().getCurrentHighlight() diff --git a/src/main/java/ghidrathon/GhidrathonConsoleInputThread.java b/src/main/java/ghidrathon/GhidrathonConsoleInputThread.java index b744e42..aae8fc7 100644 --- a/src/main/java/ghidrathon/GhidrathonConsoleInputThread.java +++ b/src/main/java/ghidrathon/GhidrathonConsoleInputThread.java @@ -26,13 +26,95 @@ import java.io.PrintWriter; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.*; + public class GhidrathonConsoleInputThread extends Thread { + public static interface GhidrathonREPLCallable { + public abstract T call(GhidrathonInterpreter python); + } + + public static class GhidrathonREPLThread extends Thread { + private static int generationCount = 0; + + private GhidrathonConfig config; + GhidrathonInterpreter python = null; + + private AtomicBoolean shouldContinue = new AtomicBoolean(true); + private LinkedBlockingQueue> tasks = new LinkedBlockingQueue<>(); + + public GhidrathonREPLThread(GhidrathonConfig config) { + super("Ghidrathon console repl thread (generation " + ++generationCount + ")"); + this.config = config; + } + + private void abort() { + shouldContinue.set(false); + + while (!tasks.isEmpty()) { + try { + tasks.take().cancel(true); + } catch (InterruptedException e) {} + } + } + + @Override + public void run() { + try { + python = GhidrathonInterpreter.get(config); + } catch (RuntimeException e) { + if (python != null) { + python.close(); + } + + abort(); + e.printStackTrace(config.getStdErr()); + return; + } + + while (shouldContinue.get() || !tasks.isEmpty()) { + try { + tasks.take().run(); + } catch (InterruptedException e) { + e.printStackTrace(config.getStdErr()); + } + } + python.close(); + } + + public Future submitTask(final GhidrathonREPLCallable f) throws InterruptedException { + if (!shouldContinue.get()) { + throw new IllegalStateException("Cannot submit task. REPL thread was already disposed"); + } + + final GhidrathonREPLThread self = this; + + RunnableFuture result = new FutureTask(() -> f.call(self.python)); + tasks.put(result); + + return result; + } + + public void dispose() { + // This order prevents race conditions + // We set shouldContinue to false first and then wake up the thread + // This ensures that when the thread processes + shouldContinue.set(false); + + // no-op future to wake up the thread + try { + tasks.put(new FutureTask(() -> null)); + } catch (InterruptedException e) { + // This shouldn't happen, but there is nothing we can do except log the error + e.printStackTrace(config.getStdErr()); + } + } + } private static int generationCount = 0; private GhidrathonPlugin plugin = null; private InterpreterConsole console = null; - private GhidrathonInterpreter python = null; + private GhidrathonREPLThread repl = null; private AtomicBoolean shouldContinue = new AtomicBoolean(true); private GhidrathonConfig config = GhidrathonUtils.getDefaultGhidrathonConfig(); @@ -49,6 +131,10 @@ public class GhidrathonConsoleInputThread extends Thread { config.addStdOut(console.getOutWriter()); } + public GhidrathonREPLThread getREPL() { + return repl; + } + /** * Console input thread. * @@ -61,26 +147,28 @@ public void run() { console.clear(); - try { - - python = GhidrathonInterpreter.get(config); + repl = new GhidrathonREPLThread(config); + repl.start(); - python.printWelcome(); - - } catch (RuntimeException e) { + try { - if (python != null) { - python.close(); - } + repl.submitTask((python) -> { + python.printWelcome(); + return null; + }).get(); + } catch (InterruptedException | ExecutionException e) { e.printStackTrace(config.getStdErr()); + Msg.error(GhidrathonConsoleInputThread.class, "Failed to start ghidrathon console.", e); + + repl.dispose(); return; } try (BufferedReader reader = new BufferedReader(new InputStreamReader(console.getStdin()))) { plugin.flushConsole(); - console.setPrompt(python.getPrimaryPrompt()); + console.setPrompt(repl.python.getPrimaryPrompt()); // begin reading and passing input from console stdin to Python to be evaluated while (shouldContinue.get()) { @@ -105,10 +193,14 @@ public void run() { this.plugin.flushConsole(); this.console.setPrompt( - moreInputWanted ? python.getSecondaryPrompt() : python.getPrimaryPrompt()); + moreInputWanted ? repl.python.getSecondaryPrompt() : repl.python.getPrimaryPrompt()); } - } catch (RuntimeException | IOException e) { + } catch (InterruptedException | ExecutionException e) { + e.printStackTrace(config.getStdErr()); + + Msg.error(GhidrathonConsoleInputThread.class, "Failed to evaluate python. Please reset", e); + } catch (IOException e) { e.printStackTrace(); Msg.error( @@ -118,7 +210,7 @@ public void run() { } finally { - python.close(); + repl.dispose(); } } @@ -132,7 +224,7 @@ public void run() { * @return True if more input needed, otherwise False * @throws RuntimeException */ - private boolean evalPython(String line) throws RuntimeException { + private boolean evalPython(String line) throws InterruptedException, ExecutionException { boolean status; @@ -157,7 +249,8 @@ private boolean evalPython(String line) throws RuntimeException { interactiveTaskMonitor, new PrintWriter(console.getStdOut())); - status = python.eval(line, interactiveScript); + status = repl.submitTask((python) -> python.eval(line, interactiveScript)).get(); + } finally { interactiveScript.end(false); } diff --git a/src/main/java/ghidrathon/GhidrathonPlugin.java b/src/main/java/ghidrathon/GhidrathonPlugin.java index e621ad9..236842a 100644 --- a/src/main/java/ghidrathon/GhidrathonPlugin.java +++ b/src/main/java/ghidrathon/GhidrathonPlugin.java @@ -24,10 +24,13 @@ import ghidra.util.task.TaskLauncher; import ghidra.util.task.TaskMonitor; import ghidra.util.task.TaskMonitorAdapter; +import ghidra.util.Msg; import java.io.OutputStream; import java.io.PrintWriter; +import java.io.File; import java.util.Collections; import java.util.List; +import java.util.concurrent.ExecutionException; import javax.swing.*; // @formatter:off @@ -113,8 +116,26 @@ public void flushConsole() { @Override public List getCompletions(String cmd) { - // TODO Auto-generated method stub - return Collections.emptyList(); + return getCompletions(cmd, cmd.length()); + } + + /** + * Returns a list of possible command completion values at the given position. + * + * @param cmd current command line (without prompt) + * @param caretPos The position of the caret in the input string 'cmd' + * @return A list of possible command completion values. Could be empty if there aren't any. + */ + @Override + public List getCompletions(String cmd, int caretPos) { + try { + return inputThread.getREPL().submitTask((python) -> python.getCommandCompletions(cmd, caretPos)).get(); + } catch (InterruptedException | ExecutionException e) { + e.printStackTrace(console.getErrWriter()); + Msg.error(GhidrathonConsoleInputThread.class, "An internal error occured while retrieving completion options. Please reset.", e); + + return Collections.emptyList(); + } } private void resetInterpreter() { diff --git a/src/main/java/ghidrathon/PythonCodeCompletionFactory.java b/src/main/java/ghidrathon/PythonCodeCompletionFactory.java new file mode 100644 index 0000000..8b2f09b --- /dev/null +++ b/src/main/java/ghidrathon/PythonCodeCompletionFactory.java @@ -0,0 +1,281 @@ +/* ### + * IP: GHIDRATHON + * + * 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. + */ +// Taken from https://github.com/NationalSecurityAgency/ghidra/blob/d7d6b44e296ac4a215766916d5c24e8b53b2909a/Ghidra/Features/Python/src/main/java/ghidra/python/PythonCodeCompletionFactory.java +package ghidrathon; + +import java.awt.Color; +import java.lang.reflect.Field; +import java.lang.reflect.Method; +import java.util.*; + +import javax.swing.JComponent; + +import jep.python.PyObject; +import jep.JepException; + +import docking.widgets.label.GDLabel; +import ghidra.app.plugin.core.console.CodeCompletion; +import ghidra.framework.options.Options; +import ghidra.util.Msg; + +/** + * Generates CodeCompletions from Python objects. + * + * + * + */ +public class PythonCodeCompletionFactory { + private static class InspectableJavaObject { + private Class srcClass; + + public InspectableJavaObject(Class c) { + srcClass = c; + } + + public Class getSrcClass() { + return srcClass; + } + + /** + * Returns the Java methods declared for a given object + * @param obj a Java Object + * @return the Java method names and a proxy that can be inspected as well + */ + public List getProperties() { + List properties = new ArrayList<>(); + Method[] declaredMethods = srcClass.getMethods(); + Field[] declaredFields = srcClass.getFields(); + + for (Method declaredMethod : declaredMethods) { + properties.add(new Object [] { + declaredMethod.getName(), + new InspectableJavaMethod(this, declaredMethod) + }); + } + + for (Field declaredField : declaredFields) { + properties.add(new Object [] { + declaredField.getName(), + new InspectableJavaObject(declaredField.getType()) + }); + } + + return properties; + } + } + + private static class InspectableJavaMethod { + private Method method; + public String __name__; + public Object __self__; + + public InspectableJavaMethod(Object o, Method m) { + method = m; + __name__ = m.getName(); + __self__ = o; + } + + public Method getMethod() { + return method; + } + } + + + private static Map typeToColorMap = new HashMap<>(); + public static final String COMPLETION_LABEL = "Code Completion Colors"; + + /* package-level accessibility so that PythonPlugin can tell this is + * our option + */ + final static String INCLUDE_TYPES_LABEL = "Include type names in code completion popup?"; + private final static String INCLUDE_TYPES_DESCRIPTION = + "Whether or not to include the type names (classes) of the possible " + + "completions in the code completion window. The class name will be " + + "parenthesized after the completion."; + private final static boolean INCLUDE_TYPES_DEFAULT = true; + private static boolean includeTypes = INCLUDE_TYPES_DEFAULT; + + public static final Color NULL_COLOR = new Color(255, 0, 0); + public static final Color FUNCTION_COLOR = new Color(0, 128, 0); + public static final Color PACKAGE_COLOR = new Color(128, 0, 0); + public static final Color CLASS_COLOR = new Color(0, 0, 255); + public static final Color METHOD_COLOR = new Color(0, 128, 128); + /* anonymous code chunks */ + public static final Color CODE_COLOR = new Color(0, 64, 0); + public static final Color INSTANCE_COLOR = new Color(128, 0, 128); + public static final Color SEQUENCE_COLOR = new Color(128, 96, 64); + public static final Color MAP_COLOR = new Color(64, 96, 128); + public static final Color NUMBER_COLOR = new Color(64, 64, 64); + + static { + setupClass("NoneType", NULL_COLOR); + + setupClass("builtin_function_or_method", FUNCTION_COLOR); + setupClass("function", FUNCTION_COLOR); + + setupClass("module", PACKAGE_COLOR); + + setupClass("str", SEQUENCE_COLOR); + setupClass("bytes", SEQUENCE_COLOR); + setupClass("list", SEQUENCE_COLOR); + setupClass("tuple", SEQUENCE_COLOR); + setupClass("dict", MAP_COLOR); + + setupClass("int", NUMBER_COLOR); + setupClass("float", NUMBER_COLOR); + setupClass("complex", NUMBER_COLOR); + } + + /** + * Returns the type name for a Python Object. + * + * @param pyObj Object to determine to type name + * @return The type name. + */ + private static String getTypeName(PyObject pyObj) { + return pyObj.getAttr("__class__", PyObject.class).getAttr("__name__", String.class); + } + + /** + * Sets up a Type mapping. + * + * @param typeName Type name + * @param defaultColor default Color for this type + * @param description description of the type + */ + private static void setupClass(String typeName, Color defaultColor) { + typeToColorMap.put(typeName, defaultColor); + } + + /** + * 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 + * @return A new CodeCompletion from the given Python objects. + */ + public static CodeCompletion newCodeCompletion(String description, String insertion, + Object obj) { + JComponent comp = null; + + if (obj != null) { + String type; + Color color; + if ((obj instanceof PyObject)) { + type = getTypeName((PyObject) obj); + color = typeToColorMap.get(type); + } else { + type = obj.getClass().getSimpleName(); + color = CLASS_COLOR; + } + if (includeTypes) { + description = description + " (" + type + ")"; + } + + comp = new GDLabel(description); + if (color != null) { + comp.setForeground(color); + } + } + return new CodeCompletion(description, insertion, comp); + } + + /** + * Returns the Java methods declared for a given object + * @param obj a Java Object + * @return the Java method names and a proxy that can be inspected as well + */ + public static Object getReturnType(Object obj, String name) { + Class c; + if (obj instanceof InspectableJavaObject) { + c = ((InspectableJavaObject) obj).getSrcClass(); + } else { + c = obj.getClass(); + } + + return new InspectableJavaObject(getReturnTypeForClass(c, name)); + } + + /** + * Returns the Java methods declared for a given java class + * @param obj a Java Object + * @return the Java method names and a proxy that can be inspected as well + */ + public static Class getReturnTypeForClass(Class c, String name) { + Method[] declaredMethods = c.getMethods(); + + for (Method declaredMethod : declaredMethods) { + if (declaredMethod.getName().equals(name)) { + return declaredMethod.getReturnType(); + } + } + + return null; + } + + /** + * Sets up Python code completion Options. + * @param plugin python plugin as options owner + * @param options an Options handle + */ + public static void setupOptions(GhidrathonPlugin plugin, Options options) { + includeTypes = options.getBoolean(INCLUDE_TYPES_LABEL, INCLUDE_TYPES_DEFAULT); + options.registerOption(INCLUDE_TYPES_LABEL, INCLUDE_TYPES_DEFAULT, null, + INCLUDE_TYPES_DESCRIPTION); + + Iterator iter = typeToColorMap.keySet().iterator(); + while (iter.hasNext()) { + String currentType = (String) iter.next(); + options.registerOption( + COMPLETION_LABEL + Options.DELIMITER + currentType, + typeToColorMap.get(currentType), null, + "Color to use for '" + currentType + "'."); + typeToColorMap.put(currentType, + options.getColor(COMPLETION_LABEL + Options.DELIMITER + currentType, + typeToColorMap.get(currentType))); + } + } + + /** + * Handle an Option change. + * + * This is named slightly differently because it is a static method, not + * an instance method. + * + * By the time we get here, we assume that the Option changed is indeed + * ours. + * + * @param options the Options handle + * @param name name of the Option changed + * @param oldValue the old value + * @param newValue the new value + */ + public static void changeOptions(Options options, String name, Object oldValue, + Object newValue) { + String typeName = name.substring((COMPLETION_LABEL + Options.DELIMITER).length()); + + if (typeToColorMap.containsKey(typeName)) { + typeToColorMap.put(typeName, (Color) newValue); + } + else if (name.equals(INCLUDE_TYPES_LABEL)) { + includeTypes = ((Boolean) newValue).booleanValue(); + } + else { + Msg.error(PythonCodeCompletionFactory.class, "unknown option '" + name + "'"); + } + } +} diff --git a/src/main/java/ghidrathon/interpreter/GhidrathonInterpreter.java b/src/main/java/ghidrathon/interpreter/GhidrathonInterpreter.java index 0a40015..9c08084 100644 --- a/src/main/java/ghidrathon/interpreter/GhidrathonInterpreter.java +++ b/src/main/java/ghidrathon/interpreter/GhidrathonInterpreter.java @@ -11,9 +11,11 @@ package ghidrathon.interpreter; import generic.jar.ResourceFile; +import ghidra.app.plugin.core.console.CodeCompletion; import ghidra.app.script.GhidraScript; import ghidra.app.script.GhidraScriptUtil; import ghidra.framework.Application; +import ghidra.util.Msg; import ghidrathon.GhidrathonClassEnquirer; import ghidrathon.GhidrathonConfig; import ghidrathon.GhidrathonScript; @@ -24,10 +26,17 @@ import java.io.PrintWriter; import java.lang.reflect.*; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.ArrayList; +import java.util.Collections; +import java.util.*; import jep.Jep; import jep.JepConfig; import jep.JepException; import jep.MainInterpreter; +import jep.python.PyObject; +import jep.python.PyCallable; + +import util.CollectionUtils; /** Utility class used to configure a Jep instance to access Ghidra */ public class GhidrathonInterpreter { @@ -37,6 +46,8 @@ public class GhidrathonInterpreter { private PrintWriter err = null; private GhidrathonConfig config = null; + private PyObject introspectModule = null; + // these variables set across GhidrathonInterpreter instances private static final JepConfig jepConfig = new JepConfig(); private static final GhidrathonClassEnquirer ghidrathonClassEnquirer = @@ -71,6 +82,9 @@ private GhidrathonInterpreter(GhidrathonConfig config) throws JepException, IOEx // create new Jep SharedInterpreter instance jep_ = new jep.SharedInterpreter(); + jep_.eval("import complete"); + introspectModule = jep_.getValue("complete", PyObject.class); + // now that everything is configured, we should be able to run some utility scripts // to help us further configure the Python environment setJepWrappers(); @@ -409,4 +423,30 @@ public String getSecondaryPrompt() { return "... "; } + + // Taken from https://github.com/NationalSecurityAgency/ghidra/blob/d7d6b44e296ac4a215766916d5c24e8b53b2909a/Ghidra/Features/Python/src/main/java/ghidra/python/GhidraPythonInterpreter.java#L440 + // Heavily modified + /** + * Returns the possible command completions for a command. + * + * @param cmd The command line. + * @param caretPos The position of the caret in the input string 'cmd' + * @return A list of possible command completions. Could be empty if there aren't any. + * @see PythonPlugin#getCompletions + */ + public List getCommandCompletions(String cmd, int caretPos) { + // At this point the caret is assumed to be positioned right after the value we need to + // complete (example: "[complete.Me, rest, code]"). To make the completion work + // in our case, it's sufficient (albeit naive) to just remove the text on the right side + // of our caret. The later code (on the python's side) will parse the rest properly + // and will generate the completions. + cmd = cmd.substring(0, caretPos); + + return (List) introspectModule.getAttr("complete", PyCallable.class).call(cmd, getLocals()); + } + + private PyObject getLocals() { + jep_.eval("__ghidrathon_locals__ = locals()"); + return jep_.getValue("__ghidrathon_locals__", PyObject.class); + } }