Skip to content

Commit 19051d9

Browse files
rbranemesare
rbran
authored and
emesare
committed
Implement PluginCommand for rust
1 parent 9262278 commit 19051d9

File tree

3 files changed

+334
-0
lines changed

3 files changed

+334
-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: 267 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,267 @@
1+
use std::ffi::{c_char, c_void, CStr};
2+
3+
use binaryninjacore_sys::*;
4+
5+
use crate::architecture::CoreArchitecture;
6+
use crate::high_level_il::HighLevelILFunction;
7+
use crate::low_level_il::RegularLowLevelILFunction;
8+
use crate::medium_level_il::MediumLevelILFunction;
9+
use crate::rc::{Array, CoreArrayProvider, CoreArrayProviderInner, Ref};
10+
use crate::string::BnStrCompatible;
11+
use crate::{BinaryView, Function};
12+
13+
pub type PluginCommandType = BNPluginCommandType;
14+
15+
#[repr(transparent)]
16+
#[derive(Debug)]
17+
pub struct PluginCommand<T> {
18+
handle: BNPluginCommand,
19+
_kind: core::marker::PhantomData<T>,
20+
}
21+
22+
impl<T> PluginCommand<T> {
23+
pub fn ty(&self) -> PluginCommandType {
24+
self.handle.type_
25+
}
26+
27+
pub fn name(&self) -> &str {
28+
unsafe { CStr::from_ptr(self.handle.name) }
29+
.to_str()
30+
.unwrap()
31+
}
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+
$($get_all:ident, $register_name:ident, $enum_name:ident, $trait_name:ident {
63+
$fun_run:ident :: $fun_is_valid:ident(
64+
$(
65+
$arg_name:ident:
66+
$raw_arg_type:ty:
67+
$arg_type:ty =
68+
$rust2ffi:expr =>
69+
$ffi2rust:expr
70+
),* $(,)?
71+
) $(,)?
72+
}),* $(,)?
73+
) => {
74+
impl PluginCommand<PluginCommandAll> {
75+
pub fn valid_plugin_commands() -> Array<Self> {
76+
let mut count = 0;
77+
let array = unsafe { BNGetAllPluginCommands(&mut count) };
78+
unsafe { Array::new(array, count, ()) }
79+
}
80+
81+
pub fn execution(&self) -> PluginCommandAll {
82+
match self.ty() {
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 c_void,
101+
is_valid: Option<unsafe extern "C" fn (
102+
ctxt: *mut c_void,
103+
view: *mut BNBinaryView,
104+
$($raw_arg_type),*
105+
) -> bool>,
106+
run: unsafe extern "C" fn (
107+
ctxt: *mut c_void,
108+
view: *mut BNBinaryView,
109+
$($raw_arg_type),*
110+
),
111+
}
112+
113+
impl $enum_name {
114+
pub fn is_valid(
115+
&mut self,
116+
view: &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, view.handle, $($rust2ffi),*) }
124+
}
125+
126+
pub fn run(
127+
&mut self,
128+
view: &BinaryView,
129+
$($arg_name: $arg_type),*
130+
) {
131+
unsafe{ (self.run)(self.context, view.handle, $($rust2ffi),*) }
132+
}
133+
}
134+
)*
135+
136+
$(
137+
pub trait $trait_name: Send + Sync + 'static {
138+
fn is_valid(
139+
&mut self,
140+
view: &BinaryView,
141+
$($arg_name: $arg_type),*
142+
) -> bool;
143+
144+
fn run(
145+
&mut self,
146+
view: &BinaryView,
147+
$($arg_name: $arg_type),*
148+
);
149+
150+
fn register(self, name: impl BnStrCompatible, description: impl BnStrCompatible) where Self: Sized {
151+
unsafe extern "C" fn ffi_action<T: $trait_name>(
152+
ctxt: *mut c_void,
153+
view: *mut BNBinaryView,
154+
$($arg_name: $raw_arg_type),*
155+
) {
156+
let slf = ctxt as *mut T;
157+
(*slf).run(&BinaryView::from_raw(view), $($ffi2rust),*)
158+
}
159+
unsafe extern "C" fn ffi_is_valid<T: $trait_name>(
160+
ctxt: *mut c_void,
161+
view: *mut BNBinaryView,
162+
$($arg_name: $raw_arg_type),*
163+
) -> bool {
164+
let slf = ctxt as *mut T;
165+
(*slf).is_valid(&BinaryView::from_raw(view), $($ffi2rust),*)
166+
}
167+
let name = name.into_bytes_with_nul();
168+
let description = description.into_bytes_with_nul();
169+
unsafe{ $register_name(
170+
name.as_ref().as_ptr() as *const c_char,
171+
description.as_ref().as_ptr() as *const c_char,
172+
Some(ffi_action::<Self>),
173+
Some(ffi_is_valid::<Self>),
174+
Box::leak(Box::new(self)) as *mut Self as *mut c_void,
175+
) }
176+
}
177+
}
178+
179+
impl PluginCommand<$enum_name> {
180+
pub fn valid_plugin_commands(view: &BinaryView, $($arg_name: $arg_type),*) -> Array<Self> {
181+
let mut count = 0;
182+
let array = unsafe { $get_all(view.handle, $($rust2ffi, )* &mut count) };
183+
unsafe { Array::new(array, count, ()) }
184+
}
185+
186+
pub fn execution(&self) -> $enum_name {
187+
assert_eq!(self.ty(), PluginCommandType::$enum_name);
188+
$enum_name {
189+
context: self.handle.context,
190+
is_valid: self.handle.$fun_is_valid,
191+
run: self.handle.$fun_run.unwrap(),
192+
}
193+
}
194+
}
195+
)*
196+
};
197+
}
198+
199+
define_enum! {
200+
BNGetValidPluginCommands, BNRegisterPluginCommand, DefaultPluginCommand, CustomDefaultPluginCommand {
201+
defaultCommand::defaultIsValid(),
202+
},
203+
BNGetValidPluginCommandsForAddress, BNRegisterPluginCommandForAddress, AddressPluginCommand, CustomAddressPluginCommand {
204+
addressCommand::addressIsValid(
205+
addr: u64: u64 = addr => addr,
206+
),
207+
},
208+
BNGetValidPluginCommandsForRange, BNRegisterPluginCommandForRange, RangePluginCommand, CustomRangePluginCommand {
209+
rangeCommand::rangeIsValid(
210+
addr: u64: u64 = addr => addr,
211+
len: u64: u64 = len => len,
212+
),
213+
},
214+
BNGetValidPluginCommandsForFunction, BNRegisterPluginCommandForFunction, FunctionPluginCommand, CustomFunctionPluginCommand {
215+
functionCommand::functionIsValid(
216+
func: *mut BNFunction: &Function = func.handle => &Function::from_raw(func),
217+
),
218+
},
219+
BNGetValidPluginCommandsForLowLevelILFunction, BNRegisterPluginCommandForLowLevelILFunction, LowLevelILFunctionPluginCommand, CustomLowLevelILFunctionPluginCommand {
220+
lowLevelILFunctionCommand::lowLevelILFunctionIsValid(
221+
llil: *mut BNLowLevelILFunction: &RegularLowLevelILFunction<CoreArchitecture> =
222+
llil.handle => &RegularLowLevelILFunction::from_raw(
223+
get_function_from_llil(llil).arch(),
224+
llil,
225+
),
226+
),
227+
},
228+
BNGetValidPluginCommandsForLowLevelILInstruction, BNRegisterPluginCommandForLowLevelILInstruction, LowLevelILInstructionPluginCommand, CustomLowLevelILInstructionPluginCommand {
229+
lowLevelILInstructionCommand::lowLevelILInstructionIsValid(
230+
llil: *mut BNLowLevelILFunction: &RegularLowLevelILFunction<CoreArchitecture> =
231+
llil.handle => &RegularLowLevelILFunction::from_raw(
232+
get_function_from_llil(llil).arch(),
233+
llil,
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+
}
262+
263+
// TODO merge this into the low_level_il module
264+
unsafe fn get_function_from_llil(llil: *mut BNLowLevelILFunction) -> Ref<Function> {
265+
let func = BNGetLowLevelILOwnerFunction(llil);
266+
Function::ref_from_raw(func)
267+
}

‎rust/tests/plugin_command.rs‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
use binaryninja::binary_view::BinaryView;
2+
use binaryninja::headless::Session;
3+
use binaryninja::plugin_command::{
4+
CustomDefaultPluginCommand, DefaultPluginCommand, PluginCommand, PluginCommandAll,
5+
};
6+
use std::path::PathBuf;
7+
use std::sync::atomic::{AtomicUsize, Ordering};
8+
use std::sync::Arc;
9+
10+
#[test]
11+
fn test_custom_plugin_command() {
12+
let _session = Session::new().expect("Failed to initialize session");
13+
let out_dir = env!("OUT_DIR").parse::<PathBuf>().unwrap();
14+
let view = binaryninja::load(out_dir.join("atox.obj")).expect("Failed to create view");
15+
let counter = Arc::new(AtomicUsize::new(0));
16+
17+
struct MyCommand {
18+
counter: Arc<AtomicUsize>,
19+
}
20+
21+
impl CustomDefaultPluginCommand for MyCommand {
22+
fn is_valid(&mut self, _view: &BinaryView) -> bool {
23+
true
24+
}
25+
26+
fn run(&mut self, _view: &BinaryView) {
27+
self.counter.fetch_add(1, Ordering::SeqCst);
28+
}
29+
}
30+
31+
// Register the plugin command.
32+
const PLUGIN_NAME: &str = "MyTestCommand1";
33+
let test_command = MyCommand {
34+
counter: counter.clone(),
35+
};
36+
test_command.register(
37+
PLUGIN_NAME,
38+
"Test for the plugin command custom implementation",
39+
);
40+
41+
// Execute the plugin command to verify that it's callable.
42+
let all_plugins = PluginCommand::<PluginCommandAll>::valid_plugin_commands();
43+
let my_core_plugin = all_plugins
44+
.iter()
45+
.find(|plugin| plugin.name() == PLUGIN_NAME)
46+
.unwrap();
47+
match my_core_plugin.execution() {
48+
PluginCommandAll::DefaultPluginCommand(mut cmd) => {
49+
assert!(cmd.is_valid(&view));
50+
cmd.run(&view);
51+
}
52+
_ => unreachable!(),
53+
}
54+
55+
// Get the plugin through the specific command type.
56+
let view_default_plugins = PluginCommand::<DefaultPluginCommand>::valid_plugin_commands(&view);
57+
let my_core_plugin_1 = view_default_plugins
58+
.iter()
59+
.find(|plugin| plugin.name() == PLUGIN_NAME)
60+
.unwrap();
61+
let mut my_core_plugin_cmd = my_core_plugin_1.execution();
62+
assert!(my_core_plugin_cmd.is_valid(&view));
63+
my_core_plugin_cmd.run(&view);
64+
65+
assert_eq!(counter.load(Ordering::SeqCst), 2);
66+
}

0 commit comments

Comments
 (0)