Add Checkbox input support to python api

6 files changed, 212 insertions, 1 deletions

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index d3bb4a15..6fcc1d97 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -2049,6 +2049,22 @@ namespace BinaryNinja {
 	*/
 	bool GetDirectoryNameInput(std::string& result, const std::string& prompt, const std::string& defaultName = "");
 
+	/*! Prompts the user for a checkbox input
+		\ingroup interaction
+
+		\param[out] result Reference to the integer the result will be copied to
+		\param[in] prompt Prompt for the dialog
+		\param[in] title Title for the input popup when used in UI
+		\param[in] defaultChoice Default checkbox state (0 == unchecked, 1 == checked)
+		\return Whether a checkbox input was successfully received
+	*/
+	bool GetCheckboxInput(
+		int64_t& result,
+		const std::string& prompt,
+		const std::string& title,
+		const int64_t& defaultChoice
+	);
+
 	/*! Prompts the user for a set of inputs specified in `fields` with given title.
 		The fields parameter is a list containing FieldInputFields
 
@@ -17027,6 +17043,7 @@ namespace BinaryNinja {
 		static FormInputField SaveFileName(
 		    const std::string& prompt, const std::string& ext, const std::string& defaultName = "");
 		static FormInputField DirectoryName(const std::string& prompt, const std::string& defaultName = "");
+		static FormInputField Checkbox(const std::string& prompt, const bool& defaultChoice = false);
 	};
 
 	/*!
@@ -17086,6 +17103,13 @@ namespace BinaryNinja {
 		    const std::string& defaultName = "");
 		virtual bool GetDirectoryNameInput(
 		    std::string& result, const std::string& prompt, const std::string& defaultName = "");
+		virtual bool GetCheckboxInput(
+			int64_t& result,
+			const std::string& prompt,
+			const std::string& title,
+			const int64_t
+			& defaultChoice = 0
+		);
 		virtual bool GetFormInput(std::vector<FormInputField>& fields, const std::string& title) = 0;
 
 		virtual BNMessageBoxButtonResult ShowMessageBox(const std::string& title, const std::string& text,
diff --git a/binaryninjacore.h b/binaryninjacore.h
index dc6f293d..9e72f01b 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -3037,7 +3037,8 @@ extern "C"
 		ChoiceFormField,
 		OpenFileNameFormField,
 		SaveFileNameFormField,
-		DirectoryNameFormField
+		DirectoryNameFormField,
+		CheckboxFormField
 	} BNFormInputFieldType;
 
 	typedef struct BNFormInputField
@@ -3083,6 +3084,7 @@ extern "C"
 		bool (*getSaveFileNameInput)(
 		    void* ctxt, char** result, const char* prompt, const char* ext, const char* defaultName);
 		bool (*getDirectoryNameInput)(void* ctxt, char** result, const char* prompt, const char* defaultName);
+		bool (*getCheckboxInput)(void* ctxt, int64_t* result, const char* prompt, const char* title, const int64_t* defaultChoice);
 		bool (*getFormInput)(void* ctxt, BNFormInputField* fields, size_t count, const char* title);
 		BNMessageBoxButtonResult (*showMessageBox)(
 		    void* ctxt, const char* title, const char* text, BNMessageBoxButtonSet buttons, BNMessageBoxIcon icon);
@@ -7312,6 +7314,7 @@ extern "C"
 	BINARYNINJACOREAPI bool BNGetSaveFileNameInput(
 	    char** result, const char* prompt, const char* ext, const char* defaultName);
 	BINARYNINJACOREAPI bool BNGetDirectoryNameInput(char** result, const char* prompt, const char* defaultName);
+	BINARYNINJACOREAPI bool BNGetCheckboxInput(int64_t* result, const char* prompt, const char* title, const int64_t* defaultChoice);
 	BINARYNINJACOREAPI bool BNGetFormInput(BNFormInputField* fields, size_t count, const char* title);
 	BINARYNINJACOREAPI void BNFreeFormInputResults(BNFormInputField* fields, size_t count);
 	BINARYNINJACOREAPI BNMessageBoxButtonResult BNShowMessageBox(
diff --git a/interaction.cpp b/interaction.cpp
index adc661c8..c25095bb 100644
--- a/interaction.cpp
+++ b/interaction.cpp
@@ -112,6 +112,17 @@ FormInputField FormInputField::DirectoryName(const string& prompt, const string&
 }
 
 
+FormInputField FormInputField::Checkbox(const string& prompt, const bool& defaultChoice)
+{
+	FormInputField result;
+	result.type = CheckboxFormField;
+	result.prompt = prompt;
+	result.intDefault = defaultChoice;
+	result.hasDefault = true;
+	return result;
+}
+
+
 void InteractionHandler::ShowMarkdownReport(
     Ref<BinaryView> view, const string& title, const string& contents, const string& plainText)
 {
@@ -189,6 +200,12 @@ bool InteractionHandler::GetDirectoryNameInput(string& result, const string& pro
 }
 
 
+bool InteractionHandler::GetCheckboxInput(int64_t& result, const std::string& prompt, const std::string& title, const int64_t& defaultChoice)
+{
+	return BinaryNinja::GetCheckboxInput(result, prompt, "Select an option", defaultChoice);
+}
+
+
 static void ShowPlainTextReportCallback(void* ctxt, BNBinaryView* view, const char* title, const char* contents)
 {
 	InteractionHandler* handler = (InteractionHandler*)ctxt;
@@ -310,6 +327,13 @@ static bool GetDirectoryNameInputCallback(void* ctxt, char** result, const char*
 }
 
 
+static bool GetCheckboxInputCallback(void* ctxt, int64_t* result, const char* prompt, const char* title, const int64_t* defaultChoice)
+{
+	InteractionHandler* handler = (InteractionHandler*)ctxt;
+	return handler->GetCheckboxInput(*result, prompt, title, *defaultChoice);
+}
+
+
 static bool GetFormInputCallback(void* ctxt, BNFormInputField* fieldBuf, size_t count, const char* title)
 {
 	InteractionHandler* handler = (InteractionHandler*)ctxt;
@@ -353,6 +377,9 @@ static bool GetFormInputCallback(void* ctxt, BNFormInputField* fieldBuf, size_t
 		case DirectoryNameFormField:
 			fields.push_back(FormInputField::DirectoryName(fieldBuf[i].prompt, fieldBuf[i].defaultName));
 			break;
+		case CheckboxFormField:
+			fields.push_back(FormInputField::Checkbox(fieldBuf[i].prompt, fieldBuf[i].intDefault));
+			break;
 		default:
 			fields.push_back(FormInputField::Label(fieldBuf[i].prompt));
 			break;
@@ -369,6 +396,7 @@ static bool GetFormInputCallback(void* ctxt, BNFormInputField* fieldBuf, size_t
 			case DirectoryNameFormField:
 				fields.back().stringDefault = fieldBuf[i].stringDefault;
 				break;
+			case CheckboxFormField:
 			case IntegerFormField:
 				fields.back().intDefault = fieldBuf[i].intDefault;
 				break;
@@ -399,6 +427,7 @@ static bool GetFormInputCallback(void* ctxt, BNFormInputField* fieldBuf, size_t
 		case DirectoryNameFormField:
 			fieldBuf[i].stringResult = BNAllocString(fields[i].stringResult.c_str());
 			break;
+		case CheckboxFormField:
 		case IntegerFormField:
 			fieldBuf[i].intResult = fields[i].intResult;
 			break;
@@ -460,6 +489,7 @@ void BinaryNinja::RegisterInteractionHandler(InteractionHandler* handler)
 	cb.getOpenFileNameInput = GetOpenFileNameInputCallback;
 	cb.getSaveFileNameInput = GetSaveFileNameInputCallback;
 	cb.getDirectoryNameInput = GetDirectoryNameInputCallback;
+	cb.getCheckboxInput = GetCheckboxInputCallback;
 	cb.getFormInput = GetFormInputCallback;
 	cb.showMessageBox = ShowMessageBoxCallback;
 	cb.openUrl = OpenUrlCallback;
@@ -590,6 +620,12 @@ bool BinaryNinja::GetDirectoryNameInput(string& result, const string& prompt, co
 }
 
 
+bool BinaryNinja::GetCheckboxInput(int64_t& result, const std::string& prompt, const std::string& title, const int64_t& defaultChoice)
+{
+	return BNGetCheckboxInput(&result, prompt.c_str(), title.c_str(), &defaultChoice);
+}
+
+
 bool BinaryNinja::GetFormInput(vector<FormInputField>& fields, const string& title)
 {
 	// Construct field list in core format
@@ -635,6 +671,8 @@ bool BinaryNinja::GetFormInput(vector<FormInputField>& fields, const string& tit
 			case DirectoryNameFormField:
 				fieldBuf[i].stringDefault = fields[i].stringDefault.c_str();
 				break;
+			case CheckboxFormField:
+				fieldBuf[i].intDefault = fields[i].intDefault;
 			case IntegerFormField:
 				fieldBuf[i].intDefault = fields[i].intDefault;
 				break;
@@ -678,6 +716,7 @@ bool BinaryNinja::GetFormInput(vector<FormInputField>& fields, const string& tit
 		case DirectoryNameFormField:
 			fields[i].stringResult = fieldBuf[i].stringResult;
 			break;
+		case CheckboxFormField:
 		case IntegerFormField:
 			fields[i].intResult = fieldBuf[i].intResult;
 			break;
diff --git a/python/interaction.py b/python/interaction.py
index c1c55095..fe370008 100644
--- a/python/interaction.py
+++ b/python/interaction.py
@@ -486,6 +486,56 @@ class DirectoryNameField:
 		self._result = value
 
 
+class CheckboxField:
+	"""
+	``CheckboxField`` prompts the user to choose a yes/no option in a checkbox.
+	Result is stored in self.result as a boolean value.
+
+	:param str prompt: Prompt to be presented to the user
+	:param bool default: Default state of the checkbox (False == unchecked, True == checked)
+	"""
+	def __init__(self, prompt, default):
+		self._prompt = prompt
+		self._result = None
+		self._default = default
+
+	def _fill_core_struct(self, value):
+		value.type = FormInputFieldType.CheckboxFormField
+		value.prompt = self._prompt
+		value.hasDefault = True
+		value.intDefault = 1 if self._default else 0
+
+	def _fill_core_result(self, value):
+		value.intResult = 1 if self.result else 0
+
+	def _get_result(self, value):
+		self._result = value.intResult != 0
+
+	@property
+	def prompt(self):
+		return self._prompt
+
+	@prompt.setter
+	def prompt(self, value):
+		self._prompt = value
+
+	@property
+	def result(self):
+		return self._result
+
+	@result.setter
+	def result(self, value):
+		self._result = value
+
+	@property
+	def default(self):
+		return self._default
+
+	@default.setter
+	def default(self, value):
+		self._default = value
+
+
 class InteractionHandler:
 	_interaction_handler = None
 
@@ -505,6 +555,7 @@ class InteractionHandler:
 		self._cb.getOpenFileNameInput = self._cb.getOpenFileNameInput.__class__(self._get_open_filename_input)
 		self._cb.getSaveFileNameInput = self._cb.getSaveFileNameInput.__class__(self._get_save_filename_input)
 		self._cb.getDirectoryNameInput = self._cb.getDirectoryNameInput.__class__(self._get_directory_name_input)
+		self._cb.getCheckboxInput = self._cb.getCheckboxInput.__class__(self._get_checkbox_input)
 		self._cb.getFormInput = self._cb.getFormInput.__class__(self._get_form_input)
 		self._cb.showMessageBox = self._cb.showMessageBox.__class__(self._show_message_box)
 		self._cb.openUrl = self._cb.openUrl.__class__(self._open_url)
@@ -650,6 +701,16 @@ class InteractionHandler:
 		except:
 			log_error(traceback.format_exc())
 
+	def _get_checkbox_input(self, ctxt, result, prompt, default_choice):
+		try:
+			value = self.get_checkbox_input(prompt, default_choice)
+			if value is None:
+				return False
+			result[0] = value
+			return True
+		except:
+			log_error(traceback.format_exc())
+
 	def _get_form_input(self, ctxt, fields, count, title):
 		try:
 			field_objs = []
@@ -714,6 +775,14 @@ class InteractionHandler:
 					        default=fields[i].stringDefault if fields[i].hasDefault else None
 					    )
 					)
+				elif fields[i].type == FormInputFieldType.CheckboxFormField:
+					log_error(fields[i].hasDefault)
+					field_objs.append(
+					    CheckboxField(
+					        fields[i].prompt,
+							default=fields[i].intDefault != 0 if fields[i].hasDefault else None
+					    )
+					)
 				else:
 					field_objs.append(LabelField(fields[i].prompt))
 			if not self.get_form_input(field_objs, title):
@@ -795,6 +864,9 @@ class InteractionHandler:
 	def get_directory_name_input(self, prompt, default_name):
 		return get_text_line_input(prompt, "Select Directory")
 
+	def get_checkbox_input(self, prompt, default_choice):
+		return get_checkbox_input(prompt, "Choose Option(s)", default_choice)
+
 	def get_form_input(self, fields, title):
 		return False
 
@@ -1361,6 +1433,22 @@ def get_directory_name_input(prompt: str, default_name: str = ""):
 	core.free_string(value)
 	return result.decode("utf-8")
 
+def get_checkbox_input(prompt: str, title: str, default: bool = False):
+	"""
+	``get_checkbox_input`` prompts the user for a checkbox input
+	:param prompt: String to prompt with
+	:param title: Title of the window when executed in the UI
+	:param default: Optional default state for the checkbox (false == unchecked, true == checked), False if not set.
+	:rtype: bool indicating the state of the checkbox
+	"""
+	default_state = ctypes.c_int64()
+	default_state.value = 1 if default else 0
+	value = ctypes.c_int64()
+	if not core.BNGetCheckboxInput(value, prompt, title, default_state):
+		return None
+	result = value.value
+	assert result is not None
+	return result != 0
 
 def get_form_input(fields, title):
 	"""
@@ -1382,6 +1470,7 @@ def get_form_input(fields, title):
 	OpenFileNameField     Prompt for file to open
 	SaveFileNameField     Prompt for file to save to
 	DirectoryNameField    Prompt for directory name
+	CheckboxFormField     Prompt for a checkbox
 	===================== ===================================================
 
 	This API is flexible and works both in the UI via a pop-up dialog and on the command-line.
diff --git a/rust/src/interaction/form.rs b/rust/src/interaction/form.rs
index 06799681..53549bf3 100644
--- a/rust/src/interaction/form.rs
+++ b/rust/src/interaction/form.rs
@@ -112,6 +112,11 @@ pub enum FormInputField {
         default: Option<String>,
         value: Option<String>,
     },
+    Checkbox {
+        prompt: String,
+        default: Option<bool>,
+        value: bool,
+    },
 }
 
 impl FormInputField {
@@ -130,6 +135,7 @@ impl FormInputField {
         let int_default = value.hasDefault.then_some(value.intDefault);
         let address_default = value.hasDefault.then_some(value.addressDefault);
         let index_default = value.hasDefault.then_some(value.indexDefault);
+        let bool_default = value.hasDefault.then_some(value.intResult != 0);
         let extension = raw_to_string(value.ext);
         let current_address = value.currentAddress;
         let string_result = raw_to_string(value.stringResult);
@@ -186,6 +192,11 @@ impl FormInputField {
                 default: string_default,
                 value: string_result,
             },
+            BNFormInputFieldType::CheckboxFormField => Self::Checkbox {
+                prompt,
+                default: bool_default,
+                value: value.intResult != 0,
+            },
         }
     }
 
@@ -258,6 +269,7 @@ impl FormInputField {
             FormInputField::OpenFileName { .. } => BNFormInputFieldType::OpenFileNameFormField,
             FormInputField::SaveFileName { .. } => BNFormInputFieldType::SaveFileNameFormField,
             FormInputField::DirectoryName { .. } => BNFormInputFieldType::DirectoryNameFormField,
+            FormInputField::Checkbox { .. } => BNFormInputFieldType::CheckboxFormField,
         }
     }
 
@@ -274,6 +286,7 @@ impl FormInputField {
             FormInputField::OpenFileName { prompt, .. } => Some(prompt.clone()),
             FormInputField::SaveFileName { prompt, .. } => Some(prompt.clone()),
             FormInputField::DirectoryName { prompt, .. } => Some(prompt.clone()),
+            FormInputField::Checkbox { prompt, .. } => Some(prompt.clone()),
         }
     }
 
@@ -325,6 +338,7 @@ impl FormInputField {
             FormInputField::OpenFileName { default, .. } => Some(default.is_some()),
             FormInputField::SaveFileName { default, .. } => Some(default.is_some()),
             FormInputField::DirectoryName { default, .. } => Some(default.is_some()),
+            FormInputField::Checkbox { default, .. } => Some(default.is_some()),
         }
     }
 
diff --git a/rust/src/interaction/handler.rs b/rust/src/interaction/handler.rs
index f854f612..11929c94 100644
--- a/rust/src/interaction/handler.rs
+++ b/rust/src/interaction/handler.rs
@@ -27,6 +27,7 @@ pub fn register_interaction_handler<R: InteractionHandler>(custom: R) {
         getOpenFileNameInput: Some(cb_get_open_file_name_input::<R>),
         getSaveFileNameInput: Some(cb_get_save_file_name_input::<R>),
         getDirectoryNameInput: Some(cb_get_directory_name_input::<R>),
+        getCheckboxInput: Some(cb_get_checkbox_input::<R>),
         getFormInput: Some(cb_get_form_input::<R>),
         showMessageBox: Some(cb_show_message_box::<R>),
         openUrl: Some(cb_open_url::<R>),
@@ -245,6 +246,26 @@ pub trait InteractionHandler: Sync + Send + 'static {
         form.get_field_with_name(prompt)
             .and_then(|f| f.try_value_string())
     }
+
+    fn get_checkbox_input(
+        &mut self,
+        prompt: &str,
+        title: &str,
+        default: Option<i64>,
+    ) -> Option<i64> {
+        let mut form = Form::new(title.to_owned());
+        form.add_field(FormInputField::Checkbox {
+            prompt: prompt.to_string(),
+            default: default.map(|b| b == 1),
+            value: default.map(|b| b == 1)?,
+        });
+        if !self.get_form_input(&mut form) {
+            return None;
+        }
+        form.get_field_with_name(prompt)
+            .and_then(|f| f.try_value_int())
+            .map(|b| if b != 0 { 1 } else { 0 })
+    }
 }
 
 pub struct InteractionHandlerTask {
@@ -543,6 +564,27 @@ unsafe extern "C" fn cb_get_directory_name_input<R: InteractionHandler>(
     }
 }
 
+unsafe extern "C" fn cb_get_checkbox_input<R: InteractionHandler>(
+    ctxt: *mut c_void,
+    result_ffi: *mut i64,
+    prompt: *const c_char,
+    title: *const c_char,
+    default_choice: *const i64,
+) -> bool {
+    let ctxt = ctxt as *mut R;
+    let prompt = raw_to_string(prompt).unwrap();
+    let title = raw_to_string(title).unwrap();
+    let default = (!default_choice.is_null()).then(|| *default_choice);
+    let result = (*ctxt).get_checkbox_input(&prompt, &title, default);
+    if let Some(result) = result {
+        unsafe { *result_ffi = result };
+        true
+    } else {
+        unsafe { *result_ffi = 0 };
+        false
+    }
+}
+
 unsafe extern "C" fn cb_get_form_input<R: InteractionHandler>(
     ctxt: *mut c_void,
     fields: *mut BNFormInputField,