Skip to content

Commit 2cfeffd

Browse files
rbranrbran
committed
implement rust plugin command
1 parent 05e3335 commit 2cfeffd

File tree

3 files changed

+313
-0
lines changed

3 files changed

+313
-0
lines changed

‎rust/src/lib.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,7 @@ pub mod main_thread;
6464
pub mod medium_level_il;
6565
pub mod metadata;
6666
pub mod platform;
67+
pub mod plugin_command;
6768
pub mod progress;
6869
pub mod project;
6970
pub mod rc;

‎rust/src/plugin_command.rs‎

Lines changed: 261 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,261 @@
1+
use core::ffi;
2+
use std::ffi::CStr;
3+
4+
use binaryninjacore_sys::*;
5+
6+
use crate::architecture::CoreArchitecture;
7+
use crate::binary_view::BinaryViewExt;
8+
use crate::high_level_il::HighLevelILFunction;
9+
use crate::low_level_il::function::{Finalized, LowLevelILFunction, NonSSA, RegularNonSSA};
10+
use crate::medium_level_il::MediumLevelILFunction;
11+
use crate::rc::{Array, CoreArrayProvider, CoreArrayProviderInner};
12+
use crate::string::BnStrCompatible;
13+
use crate::{BinaryView, Function};
14+
15+
type PluginCommandType = BNPluginCommandType;
16+
17+
#[repr(transparent)]
18+
#[derive(Debug)]
19+
pub struct PluginCommand<T> {
20+
handle: BNPluginCommand,
21+
_kind: core::marker::PhantomData<T>,
22+
}
23+
24+
impl<T> PluginCommand<T> {
25+
pub fn type_(&self) -> PluginCommandType {
26+
self.handle.type_
27+
}
28+
pub fn name(&self) -> &str {
29+
unsafe { CStr::from_ptr(self.handle.name) }
30+
.to_str()
31+
.unwrap()
32+
}
33+
pub fn description(&self) -> &str {
34+
unsafe { CStr::from_ptr(self.handle.description) }
35+
.to_str()
36+
.unwrap()
37+
}
38+
}
39+
40+
impl<T> CoreArrayProvider for PluginCommand<T> {
41+
type Raw = BNPluginCommand;
42+
type Context = ();
43+
type Wrapped<'a>
44+
= &'a PluginCommand<T>
45+
where
46+
T: 'a;
47+
}
48+
49+
unsafe impl<T> CoreArrayProviderInner for PluginCommand<T> {
50+
unsafe fn free(raw: *mut Self::Raw, _count: usize, _context: &Self::Context) {
51+
BNFreePluginCommandList(raw)
52+
}
53+
54+
unsafe fn wrap_raw<'a>(raw: &'a Self::Raw, _context: &'a Self::Context) -> Self::Wrapped<'a> {
55+
// SAFETY BNPluginCommand and PluginCommand are transparent
56+
core::mem::transmute::<&'a BNPluginCommand, &'a PluginCommand<T>>(raw)
57+
}
58+
}
59+
60+
macro_rules! define_enum {
61+
(
62+
$bv_name:ident,
63+
$($get_all:ident, $register_name:ident, $enum_name:ident, $trait_name:ident {
64+
$fun_run:ident :: $fun_is_valid:ident(
65+
$(
66+
$arg_name:ident:
67+
$raw_arg_type:ty:
68+
$arg_type:ty =
69+
$rust2ffi:expr =>
70+
$ffi2rust:expr
71+
),* $(,)?
72+
) $(,)?
73+
}),* $(,)?
74+
) => {
75+
impl PluginCommand<PluginCommandAll> {
76+
pub fn valid_plugin_commands() -> Array<Self> {
77+
let mut count = 0;
78+
let array = unsafe { BNGetAllPluginCommands(&mut count) };
79+
unsafe { Array::new(array, count, ()) }
80+
}
81+
pub fn execution(&self) -> PluginCommandAll {
82+
match self.type_() {
83+
$(PluginCommandType::$enum_name => {
84+
PluginCommandAll::$enum_name($enum_name {
85+
context: self.handle.context,
86+
is_valid: self.handle.$fun_is_valid,
87+
run: self.handle.$fun_run.unwrap(),
88+
})
89+
}),*
90+
}
91+
}
92+
}
93+
94+
pub enum PluginCommandAll {
95+
$($enum_name($enum_name)),*
96+
}
97+
98+
$(
99+
pub struct $enum_name {
100+
context: *mut ffi::c_void,
101+
is_valid: Option<unsafe extern "C" fn (
102+
ctxt: *mut ffi::c_void,
103+
$bv_name: *mut BNBinaryView,
104+
$($raw_arg_type),*
105+
) -> bool>,
106+
run: unsafe extern "C" fn (
107+
ctxt: *mut ffi::c_void,
108+
$bv_name: *mut BNBinaryView,
109+
$($raw_arg_type),*
110+
),
111+
}
112+
113+
impl $enum_name {
114+
pub fn is_valid(
115+
&mut self,
116+
$bv_name: &BinaryView,
117+
$($arg_name: $arg_type),*
118+
) -> bool {
119+
// TODO I'm assuming is_valid be null means it's always valid
120+
let Some(fun) = self.is_valid else {
121+
return true
122+
};
123+
unsafe{ fun(self.context, $bv_name.handle, $($rust2ffi),*) }
124+
}
125+
126+
pub fn run(
127+
&mut self,
128+
$bv_name: &BinaryView,
129+
$($arg_name: $arg_type),*
130+
) {
131+
unsafe{ (self.run)(self.context, $bv_name.handle, $($rust2ffi),*) }
132+
}
133+
}
134+
)*
135+
136+
$(
137+
pub trait $trait_name: Send + Sync + 'static {
138+
fn is_valid(
139+
&mut self,
140+
$bv_name: &BinaryView,
141+
$($arg_name: $arg_type),*
142+
) -> bool;
143+
fn run(
144+
&mut self,
145+
$bv_name: &BinaryView,
146+
$($arg_name: $arg_type),*
147+
);
148+
fn register(self, name: impl BnStrCompatible, description: impl BnStrCompatible) where Self: Sized {
149+
unsafe extern "C" fn ffi_action<T: $trait_name>(
150+
ctxt: *mut ffi::c_void,
151+
$bv_name: *mut BNBinaryView,
152+
$($arg_name: $raw_arg_type),*
153+
) {
154+
let slf = ctxt as *mut T;
155+
(*slf).run(&BinaryView::from_raw($bv_name), $($ffi2rust),*)
156+
}
157+
unsafe extern "C" fn ffi_is_valid<T: $trait_name>(
158+
ctxt: *mut ffi::c_void,
159+
$bv_name: *mut BNBinaryView,
160+
$($arg_name: $raw_arg_type),*
161+
) -> bool {
162+
let slf = ctxt as *mut T;
163+
(*slf).is_valid(&BinaryView::from_raw($bv_name), $($ffi2rust),*)
164+
}
165+
let name = name.into_bytes_with_nul();
166+
let description = description.into_bytes_with_nul();
167+
unsafe{ $register_name(
168+
name.as_ref().as_ptr() as *const ffi::c_char,
169+
description.as_ref().as_ptr() as *const ffi::c_char,
170+
Some(ffi_action::<Self>),
171+
Some(ffi_is_valid::<Self>),
172+
Box::leak(Box::new(self)) as *mut Self as *mut ffi::c_void,
173+
) }
174+
}
175+
}
176+
177+
impl PluginCommand<$enum_name> {
178+
pub fn valid_plugin_commands(view: &BinaryView, $($arg_name: $arg_type),*) -> Array<Self> {
179+
let mut count = 0;
180+
let array = unsafe { $get_all(view.handle, $($rust2ffi, )* &mut count) };
181+
unsafe { Array::new(array, count, ()) }
182+
}
183+
184+
pub fn execution(&self) -> $enum_name {
185+
assert_eq!(self.type_(), PluginCommandType::$enum_name);
186+
$enum_name {
187+
context: self.handle.context,
188+
is_valid: self.handle.$fun_is_valid,
189+
run: self.handle.$fun_run.unwrap(),
190+
}
191+
}
192+
}
193+
)*
194+
};
195+
}
196+
197+
define_enum! {
198+
view,
199+
BNGetValidPluginCommands, BNRegisterPluginCommand, DefaultPluginCommand, CustomDefaultPluginCommand {
200+
defaultCommand::defaultIsValid(),
201+
},
202+
BNGetValidPluginCommandsForAddress, BNRegisterPluginCommandForAddress, AddressPluginCommand, CustomAddressPluginCommand {
203+
addressCommand::addressIsValid(
204+
addr: u64: u64 = addr => addr,
205+
),
206+
},
207+
BNGetValidPluginCommandsForRange, BNRegisterPluginCommandForRange, RangePluginCommand, CustomRangePluginCommand {
208+
rangeCommand::rangeIsValid(
209+
addr: u64: u64 = addr => addr,
210+
len: u64: u64 = len => len,
211+
),
212+
},
213+
BNGetValidPluginCommandsForFunction, BNRegisterPluginCommandForFunction, FunctionPluginCommand, CustomFunctionPluginCommand {
214+
functionCommand::functionIsValid(
215+
func: *mut BNFunction: &Function = func.handle => &Function::from_raw(func),
216+
),
217+
},
218+
BNGetValidPluginCommandsForLowLevelILFunction, BNRegisterPluginCommandForLowLevelILFunction, LowLevelILFunctionPluginCommand, CustomLowLevelILFunctionPluginCommand {
219+
lowLevelILFunctionCommand::lowLevelILFunctionIsValid(
220+
// TODO I don't know what Generics should be used here
221+
func: *mut BNLowLevelILFunction: &LowLevelILFunction<CoreArchitecture, Finalized, NonSSA<RegularNonSSA>> =
222+
func.handle => &LowLevelILFunction::from_raw(
223+
BinaryView::from_raw(view).default_arch().unwrap(),
224+
func,
225+
),
226+
),
227+
},
228+
BNGetValidPluginCommandsForLowLevelILInstruction, BNRegisterPluginCommandForLowLevelILInstruction, LowLevelILInstructionPluginCommand, CustomLowLevelILInstructionPluginCommand {
229+
lowLevelILInstructionCommand::lowLevelILInstructionIsValid(
230+
func: *mut BNLowLevelILFunction: &LowLevelILFunction<CoreArchitecture, Finalized, NonSSA<RegularNonSSA>> =
231+
func.handle => &LowLevelILFunction::from_raw(
232+
BinaryView::from_raw(view).default_arch().unwrap(),
233+
func,
234+
),
235+
instr: usize: usize = instr => instr,
236+
),
237+
},
238+
BNGetValidPluginCommandsForMediumLevelILFunction, BNRegisterPluginCommandForMediumLevelILFunction, MediumLevelILFunctionPluginCommand, CustomMediumLevelILFunctionPluginCommand {
239+
mediumLevelILFunctionCommand::mediumLevelILFunctionIsValid(
240+
func: *mut BNMediumLevelILFunction: &MediumLevelILFunction = func.handle => &MediumLevelILFunction::from_raw(func),
241+
),
242+
},
243+
BNGetValidPluginCommandsForMediumLevelILInstruction, BNRegisterPluginCommandForMediumLevelILInstruction, MediumLevelILInstructionPluginCommand, CustomMediumLevelILInstructionPluginCommand {
244+
mediumLevelILInstructionCommand::mediumLevelILInstructionIsValid(
245+
func: *mut BNMediumLevelILFunction: &MediumLevelILFunction = func.handle => &MediumLevelILFunction::from_raw(func),
246+
instr: usize: usize = instr => instr,
247+
),
248+
},
249+
BNGetValidPluginCommandsForHighLevelILFunction, BNRegisterPluginCommandForHighLevelILFunction, HighLevelILFunctionPluginCommand, CustomHighLevelILFunctionPluginCommand {
250+
// TODO I don't know if the value is full_ast or not
251+
highLevelILFunctionCommand::highLevelILFunctionIsValid(
252+
func: *mut BNHighLevelILFunction: &HighLevelILFunction = func.handle => &HighLevelILFunction{full_ast: false, handle: func},
253+
),
254+
},
255+
BNGetValidPluginCommandsForHighLevelILInstruction, BNRegisterPluginCommandForHighLevelILInstruction, HighLevelILInstructionPluginCommand, CustomHighLevelILInstructionPluginCommand {
256+
highLevelILInstructionCommand::highLevelILInstructionIsValid(
257+
func: *mut BNHighLevelILFunction: &HighLevelILFunction = func.handle => &HighLevelILFunction{full_ast: false, handle: func},
258+
instr: usize: usize = instr => instr,
259+
),
260+
},
261+
}

‎rust/tests/plugin_command.rs‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
use std::path::PathBuf;
2+
use std::sync::{Arc, Mutex};
3+
4+
use binaryninja::binary_view::BinaryView;
5+
use binaryninja::headless::Session;
6+
use binaryninja::plugin_command::{CustomDefaultPluginCommand, PluginCommand, PluginCommandAll};
7+
8+
#[test]
9+
fn test_custom_plugin_command() {
10+
let _session = Session::new().expect("Failed to initialize session");
11+
let out_dir = env!("OUT_DIR").parse::<PathBuf>().unwrap();
12+
let view = binaryninja::load(out_dir.join("atox.obj")).expect("Failed to create view");
13+
let counter = Arc::new(Mutex::new(0));
14+
struct MyCommand {
15+
counter: Arc<Mutex<usize>>,
16+
}
17+
impl CustomDefaultPluginCommand for MyCommand {
18+
fn is_valid(&mut self, _view: &BinaryView) -> bool {
19+
true
20+
}
21+
22+
fn run(&mut self, _view: &BinaryView) {
23+
let mut counter = self.counter.lock().unwrap();
24+
*counter += 1;
25+
}
26+
}
27+
const PLUGIN_NAME: &str = "MyTestCommand1";
28+
MyCommand {
29+
counter: Arc::clone(&counter),
30+
}
31+
.register(
32+
PLUGIN_NAME,
33+
"Test for the plugin command custom implementation",
34+
);
35+
36+
let all_plugins = PluginCommand::<PluginCommandAll>::valid_plugin_commands();
37+
let my_core_plugin = all_plugins
38+
.iter()
39+
.find(|plugin| plugin.name() == PLUGIN_NAME)
40+
.unwrap();
41+
match my_core_plugin.execution() {
42+
PluginCommandAll::DefaultPluginCommand(mut exe) => {
43+
assert!(exe.is_valid(&view));
44+
exe.run(&view);
45+
}
46+
_ => unreachable!(),
47+
}
48+
49+
let counter = *counter.lock().unwrap();
50+
assert_eq!(counter, 1);
51+
}

0 commit comments

Comments
 (0)