Input completion through python

5 files changed, 161 insertions, 20 deletions

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index 6b8d5fc2..1e5ab511 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -4099,6 +4099,7 @@ namespace BinaryNinja
 		static void SetCurrentBasicBlockCallback(void* ctxt, BNBasicBlock* block);
 		static void SetCurrentAddressCallback(void* ctxt, uint64_t addr);
 		static void SetCurrentSelectionCallback(void* ctxt, uint64_t begin, uint64_t end);
+		static char* CompleteInputCallback(void* ctxt, const char* text, uint64_t state);
 
 		virtual void DestroyInstance();
 
@@ -4110,6 +4111,7 @@ namespace BinaryNinja
 		virtual void SetCurrentBasicBlock(BasicBlock* block);
 		virtual void SetCurrentAddress(uint64_t addr);
 		virtual void SetCurrentSelection(uint64_t begin, uint64_t end);
+		virtual std::string CompleteInput(const std::string& text, uint64_t state);
 
 		void Output(const std::string& text);
 		void Error(const std::string& text);
@@ -4118,6 +4120,9 @@ namespace BinaryNinja
 
 		void RegisterOutputListener(ScriptingOutputListener* listener);
 		void UnregisterOutputListener(ScriptingOutputListener* listener);
+
+		std::string GetDelimiters();
+		void SetDelimiters(const std::string& delimiters);
 	};
 
 	class CoreScriptingInstance: public ScriptingInstance
@@ -4133,6 +4138,7 @@ namespace BinaryNinja
 		virtual void SetCurrentBasicBlock(BasicBlock* block) override;
 		virtual void SetCurrentAddress(uint64_t addr) override;
 		virtual void SetCurrentSelection(uint64_t begin, uint64_t end) override;
+		virtual std::string CompleteInput(const std::string& text, uint64_t state) override;
 	};
 
 	class ScriptingProvider: public StaticCoreRefCountObject<BNScriptingProvider>
diff --git a/binaryninjacore.h b/binaryninjacore.h
index 441e36b7..0481c5f4 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -1741,6 +1741,7 @@ extern "C"
 		void (*setCurrentBasicBlock)(void* ctxt, BNBasicBlock* block);
 		void (*setCurrentAddress)(void* ctxt, uint64_t addr);
 		void (*setCurrentSelection)(void* ctxt, uint64_t begin, uint64_t end);
+		char* (*completeInput)(void* ctxt, const char* text, uint64_t state);
 	};
 
 	struct BNScriptingProviderCallbacks
@@ -3640,6 +3641,9 @@ extern "C"
 	BINARYNINJACOREAPI void BNUnregisterScriptingInstanceOutputListener(BNScriptingInstance* instance,
 		BNScriptingOutputListener* callbacks);
 
+	BINARYNINJACOREAPI const char* BNGetScriptingInstanceDelimiters(BNScriptingInstance* instance);
+	BINARYNINJACOREAPI void BNSetScriptingInstanceDelimiters(BNScriptingInstance* instance, const char* delimiters);
+
 	BINARYNINJACOREAPI BNScriptingProviderInputReadyState BNGetScriptingInstanceInputReadyState(
 		BNScriptingInstance* instance);
 	BINARYNINJACOREAPI BNScriptingProviderExecuteResult BNExecuteScriptInput(BNScriptingInstance* instance,
@@ -3651,6 +3655,7 @@ extern "C"
 	BINARYNINJACOREAPI void BNSetScriptingInstanceCurrentAddress(BNScriptingInstance* instance, uint64_t addr);
 	BINARYNINJACOREAPI void BNSetScriptingInstanceCurrentSelection(BNScriptingInstance* instance,
 		uint64_t begin, uint64_t end);
+	BINARYNINJACOREAPI char* BNScriptingInstanceCompleteInput(BNScriptingInstance* instance, const char* text, uint64_t state);
 
 	// Main thread actions
 	BINARYNINJACOREAPI void BNRegisterMainThread(BNMainThreadCallbacks* callbacks);
diff --git a/python/scriptingprovider.py b/python/scriptingprovider.py
index fcbca6dc..b1405790 100644
--- a/python/scriptingprovider.py
+++ b/python/scriptingprovider.py
@@ -25,6 +25,8 @@ import ctypes
 import threading
 import abc
 import sys
+import rlcompleter
+import re
 
 # Binary Ninja components
 import binaryninja
@@ -113,7 +115,10 @@ class ScriptingInstance(object):
 			self._cb.setCurrentBasicBlock = self._cb.setCurrentBasicBlock.__class__(self._set_current_basic_block)
 			self._cb.setCurrentAddress = self._cb.setCurrentAddress.__class__(self._set_current_address)
 			self._cb.setCurrentSelection = self._cb.setCurrentSelection.__class__(self._set_current_selection)
+			self._cb.completeInput = self._cb.completeInput.__class__(self._complete_input)
+			self._cb.completeInput.restype = ctypes.c_void_p
 			self.handle = core.BNInitScriptingInstance(provider.handle, self._cb)
+			self.delimiters = ' \t\n`~!@#$%^&*()-=+[{]}\\|;:\'",<>/?'
 		else:
 			self.handle = core.handle_of_type(handle, core.BNScriptingInstance)
 		self.listeners = []
@@ -189,6 +194,15 @@ class ScriptingInstance(object):
 		except:
 			log.log_error(traceback.format_exc())
 
+	def _complete_input(self, ctxt, text, state):
+		try:
+			if not isinstance(text, str):
+				text = text.decode("charmap")
+			return ctypes.cast(binaryninja.cstr(self.perform_complete_input(text, state)), ctypes.c_void_p).value
+		except:
+			log.log_error(traceback.format_exc())
+			return ctypes.cast(binaryninja.cstr(""), ctypes.c_void_p).value
+
 	@abc.abstractmethod
 	def perform_destroy_instance(self):
 		raise NotImplementedError
@@ -221,6 +235,10 @@ class ScriptingInstance(object):
 	def perform_set_current_selection(self, begin, end):
 		raise NotImplementedError
 
+	@abc.abstractmethod
+	def perform_complete_input(self, text, state):
+		raise NotImplementedError
+
 	@property
 	def input_ready_state(self):
 		return core.BNGetScriptingInstanceInputReadyState(self.handle)
@@ -262,6 +280,9 @@ class ScriptingInstance(object):
 	def set_current_selection(self, begin, end):
 		core.BNSetScriptingInstanceCurrentSelection(self.handle, begin, end)
 
+	def complete_input(self, text, state):
+		return core.BNScriptingInstanceCompleteInput(self.handlue, text, state)
+
 	def register_output_listener(self, listener):
 		listener._register(self.handle)
 		self.listeners.append(listener)
@@ -270,6 +291,14 @@ class ScriptingInstance(object):
 		if listener in self.listeners:
 			listener._unregister(self.handle)
 			self.listeners.remove(listener)
+	
+	@property
+	def delimiters(self):
+		return core.BNGetScriptingInstanceDelimiters(self.handle)
+
+	@delimiters.setter
+	def delimiters(self, value):
+		core.BNSetScriptingInstanceDelimiters(self.handle, value)
 
 
 class _ScriptingProviderMetaclass(type):
@@ -507,6 +536,8 @@ class PythonScriptingInstance(ScriptingInstance):
 			self.code = None
 			self.input = ""
 
+			self.completer = rlcompleter.Completer(namespace = self.locals)
+
 			self.interpreter.push("from binaryninja import *")
 
 		def execute(self, code):
@@ -549,26 +580,7 @@ class PythonScriptingInstance(ScriptingInstance):
 
 					PythonScriptingInstance._interpreter.value = self
 					try:
-						self.active_view = self.current_view
-						self.active_func = self.current_func
-						self.active_block = self.current_block
-						self.active_addr = self.current_addr
-						self.active_selection_begin = self.current_selection_begin
-						self.active_selection_end = self.current_selection_end
-
-						self.locals["current_view"] = self.active_view
-						self.locals["bv"] = self.active_view
-						self.locals["current_function"] = self.active_func
-						self.locals["current_basic_block"] = self.active_block
-						self.locals["current_address"] = self.active_addr
-						self.locals["here"] = self.active_addr
-						self.locals["current_selection"] = (self.active_selection_begin, self.active_selection_end)
-						if self.active_func is None:
-							self.locals["current_llil"] = None
-							self.locals["current_mlil"] = None
-						else:
-							self.locals["current_llil"] = self.active_func.llil
-							self.locals["current_mlil"] = self.active_func.mlil
+						self.update_locals()
 
 						for line in code.split(b'\n'):
 							self.interpreter.push(line.decode('charmap'))
@@ -595,6 +607,29 @@ class PythonScriptingInstance(ScriptingInstance):
 						PythonScriptingInstance._interpreter.value = None
 						self.instance.input_ready_state = ScriptingProviderInputReadyState.ReadyForScriptExecution
 
+		def update_locals(self):
+			self.active_view = self.current_view
+			self.active_func = self.current_func
+			self.active_block = self.current_block
+			self.active_addr = self.current_addr
+			self.active_selection_begin = self.current_selection_begin
+			self.active_selection_end = self.current_selection_end
+
+			self.locals["current_view"] = self.active_view
+			self.locals["bv"] = self.active_view
+			self.locals["current_function"] = self.active_func
+			self.locals["current_basic_block"] = self.active_block
+			self.locals["current_address"] = self.active_addr
+			self.locals["here"] = self.active_addr
+			self.locals["current_selection"] = (self.active_selection_begin, self.active_selection_end)
+			if self.active_func is None:
+				self.locals["current_llil"] = None
+				self.locals["current_mlil"] = None
+			else:
+				self.locals["current_llil"] = self.active_func.llil
+				self.locals["current_mlil"] = self.active_func.mlil
+
+
 		def get_selected_data(self):
 			if self.active_view is None:
 				return None
@@ -680,6 +715,13 @@ class PythonScriptingInstance(ScriptingInstance):
 		self.interpreter.current_selection_begin = begin
 		self.interpreter.current_selection_end = end
 
+	@abc.abstractmethod
+	def perform_complete_input(self, text, state):
+		self.interpreter.update_locals()
+		result = self.interpreter.completer.complete(text, state)
+		if result is None:
+			return ""
+		return result
 
 class PythonScriptingProvider(ScriptingProvider):
 	name = "Python"
diff --git a/scriptingprovider.cpp b/scriptingprovider.cpp
index 5e6d1a02..ed010100 100644
--- a/scriptingprovider.cpp
+++ b/scriptingprovider.cpp
@@ -61,6 +61,7 @@ ScriptingInstance::ScriptingInstance(ScriptingProvider* provider)
 	cb.setCurrentBasicBlock = SetCurrentBasicBlockCallback;
 	cb.setCurrentAddress = SetCurrentAddressCallback;
 	cb.setCurrentSelection = SetCurrentSelectionCallback;
+	cb.completeInput = CompleteInputCallback;
 	AddRefForRegistration();
 	m_object = BNInitScriptingInstance(provider->GetObject(), &cb);
 }
@@ -128,6 +129,13 @@ void ScriptingInstance::SetCurrentSelectionCallback(void* ctxt, uint64_t begin,
 }
 
 
+char* ScriptingInstance::CompleteInputCallback(void* ctxt, const char* text, uint64_t state)
+{
+	ScriptingInstance* instance = (ScriptingInstance*)ctxt;
+	return BNAllocString(instance->CompleteInput(text, state).c_str());
+}
+
+
 void ScriptingInstance::DestroyInstance()
 {
 	ReleaseForRegistration();
@@ -163,6 +171,12 @@ void ScriptingInstance::SetCurrentSelection(uint64_t, uint64_t)
 }
 
 
+std::string ScriptingInstance::CompleteInput(const std::string&, uint64_t)
+{
+	return "";
+}
+
+
 void ScriptingInstance::Output(const string& text)
 {
 	BNNotifyOutputForScriptingInstance(m_object, text.c_str());
@@ -199,6 +213,18 @@ void ScriptingInstance::UnregisterOutputListener(ScriptingOutputListener* listen
 }
 
 
+std::string ScriptingInstance::GetDelimiters()
+{
+	return BNGetScriptingInstanceDelimiters(m_object);
+}
+
+
+void ScriptingInstance::SetDelimiters(const std::string& delimiters)
+{
+	BNSetScriptingInstanceDelimiters(m_object, delimiters.c_str());
+}
+
+
 CoreScriptingInstance::CoreScriptingInstance(BNScriptingInstance* instance): ScriptingInstance(instance)
 {
 }
@@ -246,6 +272,15 @@ void CoreScriptingInstance::SetCurrentSelection(uint64_t begin, uint64_t end)
 }
 
 
+std::string CoreScriptingInstance::CompleteInput(const std::string& text, uint64_t state)
+{
+	char* result = BNScriptingInstanceCompleteInput(m_object, text.c_str(), state);
+	std::string ret = result;
+	BNFreeString(result);
+	return ret;
+}
+
+
 ScriptingProvider::ScriptingProvider(const string& name): m_nameForRegister(name)
 {
 }
diff --git a/ui/scriptingconsole.h b/ui/scriptingconsole.h
index 482516c3..6c438f16 100644
--- a/ui/scriptingconsole.h
+++ b/ui/scriptingconsole.h
@@ -6,10 +6,14 @@
 #include <QtWidgets/QPushButton>
 #include <QtWidgets/QTextEdit>
 #include <QtWidgets/QWidget>
+#include <QtWidgets/QDialog>
+#include <QtWidgets/QListView>
+#include <QtCore/QAbstractListModel>
 #include <mutex>
 #include <string>
 #include <utility>
 #include <vector>
+#include <functional>
 #include "binaryninjaapi.h"
 #include "action.h"
 #include "dockhandler.h"
@@ -20,18 +24,67 @@
 
 class ScriptingConsole;
 
+class BINARYNINJAUIAPI ScriptingCompletionModel: public QAbstractListModel
+{
+	Q_OBJECT
+	std::vector<std::string> m_completions;
+
+public:
+	ScriptingCompletionModel(QWidget* parent);
+
+	virtual QModelIndex index(int row, int col, const QModelIndex& parent = QModelIndex()) const override;
+	virtual QModelIndex parent(const QModelIndex& i) const override;
+	virtual bool hasChildren(const QModelIndex& parent) const override;
+	virtual int rowCount(const QModelIndex& parent = QModelIndex()) const override;
+	virtual int columnCount(const QModelIndex& parent) const override;
+	virtual QVariant data(const QModelIndex& i, int role) const override;
+	void setModelData(const std::vector<std::string>& completions);
+};
+
+class BINARYNINJAUIAPI ScriptingCompletionPopup: public QDialog
+{
+	Q_OBJECT
+
+	QListView* m_list;
+	ScriptingCompletionModel* m_model;
+
+public:
+	ScriptingCompletionPopup(QWidget* parent);
+	void showWithData(QPoint pt, int cursorSize, const std::vector<std::string>& completions);
+	bool handleKeyEvent(QKeyEvent* event);
+
+private Q_SLOTS:
+	void clickRow(const QModelIndex& index);
+
+Q_SIGNALS:
+	void complete(QString text);
+};
+
 class BINARYNINJAUIAPI ScriptingConsoleEdit: public QTextEdit
 {
 	Q_OBJECT
+public:
+	typedef std::function<std::vector<std::string>(const std::string&)> CompletionCallback;
 
+private:
 	ScriptingConsole* m_console;
 	int m_charHeight;
 	bool m_continuation;
+	CompletionCallback m_completionCallback;
+	ScriptingCompletionPopup* m_popup;
+
+	uint64_t m_completionRegionStart;
+	uint64_t m_completionRegionInitialStop;
+	uint64_t m_completionRegionStop;
 
 public:
 	ScriptingConsoleEdit(ScriptingConsole* parent);
 	void setCharHeight(int height);
 	void setContinutation(bool cont);
+	void setCompletionCallback(CompletionCallback callback) { m_completionCallback = callback; }
+
+private Q_SLOTS:
+	void complete(QString text);
 
 protected:
 	virtual void keyPressEvent(QKeyEvent* event) override;