From 2cfeffd3f4ce585ef22378201b894dd33a6815d8 Mon Sep 17 00:00:00 2001 From: rbran Date: Mon, 7 Apr 2025 16:02:39 +0000 Subject: [PATCH 1/2] implement rust plugin command --- rust/src/lib.rs | 1 + rust/src/plugin_command.rs | 261 +++++++++++++++++++++++++++++++++++ rust/tests/plugin_command.rs | 51 +++++++ 3 files changed, 313 insertions(+) create mode 100644 rust/src/plugin_command.rs create mode 100644 rust/tests/plugin_command.rs diff --git a/rust/src/lib.rs b/rust/src/lib.rs index a3df0cd35..b2d72a4d8 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -64,6 +64,7 @@ pub mod main_thread; pub mod medium_level_il; pub mod metadata; pub mod platform; +pub mod plugin_command; pub mod progress; pub mod project; pub mod rc; diff --git a/rust/src/plugin_command.rs b/rust/src/plugin_command.rs new file mode 100644 index 000000000..bd1786e23 --- /dev/null +++ b/rust/src/plugin_command.rs @@ -0,0 +1,261 @@ +use core::ffi; +use std::ffi::CStr; + +use binaryninjacore_sys::*; + +use crate::architecture::CoreArchitecture; +use crate::binary_view::BinaryViewExt; +use crate::high_level_il::HighLevelILFunction; +use crate::low_level_il::function::{Finalized, LowLevelILFunction, NonSSA, RegularNonSSA}; +use crate::medium_level_il::MediumLevelILFunction; +use crate::rc::{Array, CoreArrayProvider, CoreArrayProviderInner}; +use crate::string::BnStrCompatible; +use crate::{BinaryView, Function}; + +type PluginCommandType = BNPluginCommandType; + +#[repr(transparent)] +#[derive(Debug)] +pub struct PluginCommand { + handle: BNPluginCommand, + _kind: core::marker::PhantomData, +} + +impl PluginCommand { + pub fn type_(&self) -> PluginCommandType { + self.handle.type_ + } + pub fn name(&self) -> &str { + unsafe { CStr::from_ptr(self.handle.name) } + .to_str() + .unwrap() + } + pub fn description(&self) -> &str { + unsafe { CStr::from_ptr(self.handle.description) } + .to_str() + .unwrap() + } +} + +impl CoreArrayProvider for PluginCommand { + type Raw = BNPluginCommand; + type Context = (); + type Wrapped<'a> + = &'a PluginCommand + where + T: 'a; +} + +unsafe impl CoreArrayProviderInner for PluginCommand { + unsafe fn free(raw: *mut Self::Raw, _count: usize, _context: &Self::Context) { + BNFreePluginCommandList(raw) + } + + unsafe fn wrap_raw<'a>(raw: &'a Self::Raw, _context: &'a Self::Context) -> Self::Wrapped<'a> { + // SAFETY BNPluginCommand and PluginCommand are transparent + core::mem::transmute::<&'a BNPluginCommand, &'a PluginCommand>(raw) + } +} + +macro_rules! define_enum { + ( + $bv_name:ident, + $($get_all:ident, $register_name:ident, $enum_name:ident, $trait_name:ident { + $fun_run:ident :: $fun_is_valid:ident( + $( + $arg_name:ident: + $raw_arg_type:ty: + $arg_type:ty = + $rust2ffi:expr => + $ffi2rust:expr + ),* $(,)? + ) $(,)? + }),* $(,)? + ) => { + impl PluginCommand { + pub fn valid_plugin_commands() -> Array { + let mut count = 0; + let array = unsafe { BNGetAllPluginCommands(&mut count) }; + unsafe { Array::new(array, count, ()) } + } + pub fn execution(&self) -> PluginCommandAll { + match self.type_() { + $(PluginCommandType::$enum_name => { + PluginCommandAll::$enum_name($enum_name { + context: self.handle.context, + is_valid: self.handle.$fun_is_valid, + run: self.handle.$fun_run.unwrap(), + }) + }),* + } + } + } + + pub enum PluginCommandAll { + $($enum_name($enum_name)),* + } + + $( + pub struct $enum_name { + context: *mut ffi::c_void, + is_valid: Option bool>, + run: unsafe extern "C" fn ( + ctxt: *mut ffi::c_void, + $bv_name: *mut BNBinaryView, + $($raw_arg_type),* + ), + } + + impl $enum_name { + pub fn is_valid( + &mut self, + $bv_name: &BinaryView, + $($arg_name: $arg_type),* + ) -> bool { + // TODO I'm assuming is_valid be null means it's always valid + let Some(fun) = self.is_valid else { + return true + }; + unsafe{ fun(self.context, $bv_name.handle, $($rust2ffi),*) } + } + + pub fn run( + &mut self, + $bv_name: &BinaryView, + $($arg_name: $arg_type),* + ) { + unsafe{ (self.run)(self.context, $bv_name.handle, $($rust2ffi),*) } + } + } + )* + + $( + pub trait $trait_name: Send + Sync + 'static { + fn is_valid( + &mut self, + $bv_name: &BinaryView, + $($arg_name: $arg_type),* + ) -> bool; + fn run( + &mut self, + $bv_name: &BinaryView, + $($arg_name: $arg_type),* + ); + fn register(self, name: impl BnStrCompatible, description: impl BnStrCompatible) where Self: Sized { + unsafe extern "C" fn ffi_action( + ctxt: *mut ffi::c_void, + $bv_name: *mut BNBinaryView, + $($arg_name: $raw_arg_type),* + ) { + let slf = ctxt as *mut T; + (*slf).run(&BinaryView::from_raw($bv_name), $($ffi2rust),*) + } + unsafe extern "C" fn ffi_is_valid( + ctxt: *mut ffi::c_void, + $bv_name: *mut BNBinaryView, + $($arg_name: $raw_arg_type),* + ) -> bool { + let slf = ctxt as *mut T; + (*slf).is_valid(&BinaryView::from_raw($bv_name), $($ffi2rust),*) + } + let name = name.into_bytes_with_nul(); + let description = description.into_bytes_with_nul(); + unsafe{ $register_name( + name.as_ref().as_ptr() as *const ffi::c_char, + description.as_ref().as_ptr() as *const ffi::c_char, + Some(ffi_action::), + Some(ffi_is_valid::), + Box::leak(Box::new(self)) as *mut Self as *mut ffi::c_void, + ) } + } + } + + impl PluginCommand<$enum_name> { + pub fn valid_plugin_commands(view: &BinaryView, $($arg_name: $arg_type),*) -> Array { + let mut count = 0; + let array = unsafe { $get_all(view.handle, $($rust2ffi, )* &mut count) }; + unsafe { Array::new(array, count, ()) } + } + + pub fn execution(&self) -> $enum_name { + assert_eq!(self.type_(), PluginCommandType::$enum_name); + $enum_name { + context: self.handle.context, + is_valid: self.handle.$fun_is_valid, + run: self.handle.$fun_run.unwrap(), + } + } + } + )* + }; +} + +define_enum! { + view, + BNGetValidPluginCommands, BNRegisterPluginCommand, DefaultPluginCommand, CustomDefaultPluginCommand { + defaultCommand::defaultIsValid(), + }, + BNGetValidPluginCommandsForAddress, BNRegisterPluginCommandForAddress, AddressPluginCommand, CustomAddressPluginCommand { + addressCommand::addressIsValid( + addr: u64: u64 = addr => addr, + ), + }, + BNGetValidPluginCommandsForRange, BNRegisterPluginCommandForRange, RangePluginCommand, CustomRangePluginCommand { + rangeCommand::rangeIsValid( + addr: u64: u64 = addr => addr, + len: u64: u64 = len => len, + ), + }, + BNGetValidPluginCommandsForFunction, BNRegisterPluginCommandForFunction, FunctionPluginCommand, CustomFunctionPluginCommand { + functionCommand::functionIsValid( + func: *mut BNFunction: &Function = func.handle => &Function::from_raw(func), + ), + }, + BNGetValidPluginCommandsForLowLevelILFunction, BNRegisterPluginCommandForLowLevelILFunction, LowLevelILFunctionPluginCommand, CustomLowLevelILFunctionPluginCommand { + lowLevelILFunctionCommand::lowLevelILFunctionIsValid( + // TODO I don't know what Generics should be used here + func: *mut BNLowLevelILFunction: &LowLevelILFunction> = + func.handle => &LowLevelILFunction::from_raw( + BinaryView::from_raw(view).default_arch().unwrap(), + func, + ), + ), + }, + BNGetValidPluginCommandsForLowLevelILInstruction, BNRegisterPluginCommandForLowLevelILInstruction, LowLevelILInstructionPluginCommand, CustomLowLevelILInstructionPluginCommand { + lowLevelILInstructionCommand::lowLevelILInstructionIsValid( + func: *mut BNLowLevelILFunction: &LowLevelILFunction> = + func.handle => &LowLevelILFunction::from_raw( + BinaryView::from_raw(view).default_arch().unwrap(), + func, + ), + instr: usize: usize = instr => instr, + ), + }, + BNGetValidPluginCommandsForMediumLevelILFunction, BNRegisterPluginCommandForMediumLevelILFunction, MediumLevelILFunctionPluginCommand, CustomMediumLevelILFunctionPluginCommand { + mediumLevelILFunctionCommand::mediumLevelILFunctionIsValid( + func: *mut BNMediumLevelILFunction: &MediumLevelILFunction = func.handle => &MediumLevelILFunction::from_raw(func), + ), + }, + BNGetValidPluginCommandsForMediumLevelILInstruction, BNRegisterPluginCommandForMediumLevelILInstruction, MediumLevelILInstructionPluginCommand, CustomMediumLevelILInstructionPluginCommand { + mediumLevelILInstructionCommand::mediumLevelILInstructionIsValid( + func: *mut BNMediumLevelILFunction: &MediumLevelILFunction = func.handle => &MediumLevelILFunction::from_raw(func), + instr: usize: usize = instr => instr, + ), + }, + BNGetValidPluginCommandsForHighLevelILFunction, BNRegisterPluginCommandForHighLevelILFunction, HighLevelILFunctionPluginCommand, CustomHighLevelILFunctionPluginCommand { + // TODO I don't know if the value is full_ast or not + highLevelILFunctionCommand::highLevelILFunctionIsValid( + func: *mut BNHighLevelILFunction: &HighLevelILFunction = func.handle => &HighLevelILFunction{full_ast: false, handle: func}, + ), + }, + BNGetValidPluginCommandsForHighLevelILInstruction, BNRegisterPluginCommandForHighLevelILInstruction, HighLevelILInstructionPluginCommand, CustomHighLevelILInstructionPluginCommand { + highLevelILInstructionCommand::highLevelILInstructionIsValid( + func: *mut BNHighLevelILFunction: &HighLevelILFunction = func.handle => &HighLevelILFunction{full_ast: false, handle: func}, + instr: usize: usize = instr => instr, + ), + }, +} diff --git a/rust/tests/plugin_command.rs b/rust/tests/plugin_command.rs new file mode 100644 index 000000000..ee2ab54ab --- /dev/null +++ b/rust/tests/plugin_command.rs @@ -0,0 +1,51 @@ +use std::path::PathBuf; +use std::sync::{Arc, Mutex}; + +use binaryninja::binary_view::BinaryView; +use binaryninja::headless::Session; +use binaryninja::plugin_command::{CustomDefaultPluginCommand, PluginCommand, PluginCommandAll}; + +#[test] +fn test_custom_plugin_command() { + let _session = Session::new().expect("Failed to initialize session"); + let out_dir = env!("OUT_DIR").parse::().unwrap(); + let view = binaryninja::load(out_dir.join("atox.obj")).expect("Failed to create view"); + let counter = Arc::new(Mutex::new(0)); + struct MyCommand { + counter: Arc>, + } + impl CustomDefaultPluginCommand for MyCommand { + fn is_valid(&mut self, _view: &BinaryView) -> bool { + true + } + + fn run(&mut self, _view: &BinaryView) { + let mut counter = self.counter.lock().unwrap(); + *counter += 1; + } + } + const PLUGIN_NAME: &str = "MyTestCommand1"; + MyCommand { + counter: Arc::clone(&counter), + } + .register( + PLUGIN_NAME, + "Test for the plugin command custom implementation", + ); + + let all_plugins = PluginCommand::::valid_plugin_commands(); + let my_core_plugin = all_plugins + .iter() + .find(|plugin| plugin.name() == PLUGIN_NAME) + .unwrap(); + match my_core_plugin.execution() { + PluginCommandAll::DefaultPluginCommand(mut exe) => { + assert!(exe.is_valid(&view)); + exe.run(&view); + } + _ => unreachable!(), + } + + let counter = *counter.lock().unwrap(); + assert_eq!(counter, 1); +} From 78cb3a4d34f8058cd07aa1bb8e838319ef7b4c57 Mon Sep 17 00:00:00 2001 From: rbran Date: Tue, 8 Apr 2025 10:21:03 +0000 Subject: [PATCH 2/2] fix rust plugin command LLIL definition --- rust/src/plugin_command.rs | 28 ++++++++++++++++------------ 1 file changed, 16 insertions(+), 12 deletions(-) diff --git a/rust/src/plugin_command.rs b/rust/src/plugin_command.rs index bd1786e23..4b3cb6585 100644 --- a/rust/src/plugin_command.rs +++ b/rust/src/plugin_command.rs @@ -4,11 +4,10 @@ use std::ffi::CStr; use binaryninjacore_sys::*; use crate::architecture::CoreArchitecture; -use crate::binary_view::BinaryViewExt; use crate::high_level_il::HighLevelILFunction; -use crate::low_level_il::function::{Finalized, LowLevelILFunction, NonSSA, RegularNonSSA}; +use crate::low_level_il::RegularLowLevelILFunction; use crate::medium_level_il::MediumLevelILFunction; -use crate::rc::{Array, CoreArrayProvider, CoreArrayProviderInner}; +use crate::rc::{Array, CoreArrayProvider, CoreArrayProviderInner, Ref}; use crate::string::BnStrCompatible; use crate::{BinaryView, Function}; @@ -217,20 +216,19 @@ define_enum! { }, BNGetValidPluginCommandsForLowLevelILFunction, BNRegisterPluginCommandForLowLevelILFunction, LowLevelILFunctionPluginCommand, CustomLowLevelILFunctionPluginCommand { lowLevelILFunctionCommand::lowLevelILFunctionIsValid( - // TODO I don't know what Generics should be used here - func: *mut BNLowLevelILFunction: &LowLevelILFunction> = - func.handle => &LowLevelILFunction::from_raw( - BinaryView::from_raw(view).default_arch().unwrap(), - func, + llil: *mut BNLowLevelILFunction: &RegularLowLevelILFunction = + llil.handle => &RegularLowLevelILFunction::from_raw( + get_function_from_llil(llil).arch(), + llil, ), ), }, BNGetValidPluginCommandsForLowLevelILInstruction, BNRegisterPluginCommandForLowLevelILInstruction, LowLevelILInstructionPluginCommand, CustomLowLevelILInstructionPluginCommand { lowLevelILInstructionCommand::lowLevelILInstructionIsValid( - func: *mut BNLowLevelILFunction: &LowLevelILFunction> = - func.handle => &LowLevelILFunction::from_raw( - BinaryView::from_raw(view).default_arch().unwrap(), - func, + llil: *mut BNLowLevelILFunction: &RegularLowLevelILFunction = + llil.handle => &RegularLowLevelILFunction::from_raw( + get_function_from_llil(llil).arch(), + llil, ), instr: usize: usize = instr => instr, ), @@ -259,3 +257,9 @@ define_enum! { ), }, } + +// TODO merge this into the low_level_il module +unsafe fn get_function_from_llil(llil: *mut BNLowLevelILFunction) -> Ref { + let func = BNGetLowLevelILOwnerFunction(llil); + Function::ref_from_raw(func) +}