Skip to content

Commit

Permalink
feat: Scanner callback API
Browse files Browse the repository at this point in the history
  • Loading branch information
Ivan Krivosheev committed Jul 21, 2021
1 parent d1e3491 commit 1b52d34
Show file tree
Hide file tree
Showing 2 changed files with 117 additions and 69 deletions.
96 changes: 41 additions & 55 deletions src/internals/scan.rs
View file
Expand Up @@ -67,22 +67,18 @@ pub fn rules_scan_mem<'a>(
flags: i32,
callback: impl FnMut(CallbackMsg<'a>) -> CallbackReturn,
) -> Result<(), YaraError> {
let p_callback: Box<Box<dyn FnMut(CallbackMsg<'a>) -> CallbackReturn>> =
Box::new(Box::new(callback));
let user_data = Box::into_raw(p_callback) as *mut c_void;
let p_callback: Box<dyn FnMut(CallbackMsg<'a>) -> CallbackReturn> = Box::new(callback);
let result = unsafe {
yara_sys::yr_rules_scan_mem(
rules,
mem.as_ptr(),
mem.len().try_into().unwrap(),
flags,
Some(scan_callback),
user_data,
&p_callback as *const Box<dyn FnMut(CallbackMsg<'a>) -> CallbackReturn> as *mut _,
timeout,
)
};
let _: Box<Box<dyn FnMut(CallbackMsg<'a>) -> CallbackReturn>> =
unsafe { Box::from_raw(user_data as *mut _) };

yara_sys::Error::from_code(result)
.map_err(|e| e.into())
Expand All @@ -98,16 +94,15 @@ pub fn scanner_scan_mem<'a>(
mem: &[u8],
callback: impl FnMut(CallbackMsg<'a>) -> CallbackReturn,
) -> Result<(), YaraError> {
let p_callback: Box<Box<dyn FnMut(CallbackMsg<'a>) -> CallbackReturn>> =
Box::new(Box::new(callback));
let user_data = Box::into_raw(p_callback) as *mut c_void;
let p_callback: Box<dyn FnMut(CallbackMsg<'a>) -> CallbackReturn> = Box::new(callback);
let result = unsafe {
yara_sys::yr_scanner_set_callback(scanner, Some(scan_callback), user_data);
yara_sys::yr_scanner_set_callback(
scanner,
Some(scan_callback),
&p_callback as *const Box<dyn FnMut(CallbackMsg<'a>) -> CallbackReturn> as *mut _,
);
yara_sys::yr_scanner_scan_mem(scanner, mem.as_ptr(), mem.len().try_into().unwrap())
};
let _: Box<Box<dyn FnMut(CallbackMsg<'a>) -> CallbackReturn>> =
unsafe { Box::from_raw(user_data as *mut _) };

yara_sys::Error::from_code(result)
.map_err(|e| e.into())
.map(|_| ())
Expand Down Expand Up @@ -152,15 +147,17 @@ pub fn rules_scan_raw<'a>(
callback: impl FnMut(CallbackMsg<'a>) -> CallbackReturn,
) -> i32 {
let fd = file.as_raw_fd();
let p_callback: Box<Box<dyn FnMut(CallbackMsg<'a>) -> CallbackReturn>> =
Box::new(Box::new(callback));
let user_data = Box::into_raw(p_callback) as *mut c_void;
let result = unsafe {
yara_sys::yr_rules_scan_fd(rules, fd, flags, Some(scan_callback), user_data, timeout)
};
let _: Box<Box<dyn FnMut(CallbackMsg<'a>) -> CallbackReturn>> =
unsafe { Box::from_raw(user_data as *mut _) };
result
let p_callback: Box<dyn FnMut(CallbackMsg<'a>) -> CallbackReturn> = Box::new(callback);
unsafe {
yara_sys::yr_rules_scan_fd(
rules,
fd,
flags,
Some(scan_callback),
&p_callback as *const Box<dyn FnMut(CallbackMsg<'a>) -> CallbackReturn> as *mut _,
timeout,
)
}
}

#[cfg(windows)]
Expand All @@ -172,23 +169,18 @@ pub fn rules_scan_raw<'a>(
callback: impl FnMut(CallbackMsg<'a>) -> CallbackReturn,
) -> i32 {
let handle = file.as_raw_handle();
let p_callback: Box<Box<dyn FnMut(CallbackMsg<'a>) -> CallbackReturn>> =
Box::new(Box::new(callback));
let user_data = Box::into_raw(p_callback) as *mut c_void;
let p_callback: Box<dyn FnMut(CallbackMsg<'a>) -> CallbackReturn> = Box::new(callback);

let result = unsafe {
unsafe {
yara_sys::yr_rules_scan_fd(
rules,
handle,
flags,
Some(scan_callback),
user_data,
&p_callback as *const Box<dyn FnMut(CallbackMsg<'a>) -> CallbackReturn> as *mut _,
timeout,
)
};
let _: Box<Box<dyn FnMut(CallbackMsg<'a>) -> CallbackReturn>> =
unsafe { Box::from_raw(user_data as *mut _) };
result
}
}

#[cfg(unix)]
Expand All @@ -202,16 +194,15 @@ pub fn scanner_scan_raw<'a>(
callback: impl FnMut(CallbackMsg<'a>) -> CallbackReturn,
) -> i32 {
let fd = file.as_raw_fd();
let p_callback: Box<Box<dyn FnMut(CallbackMsg<'a>) -> CallbackReturn>> =
Box::new(Box::new(callback));
let user_data = Box::into_raw(p_callback) as *mut c_void;
let result = unsafe {
yara_sys::yr_scanner_set_callback(scanner, Some(scan_callback), user_data);
let p_callback: Box<dyn FnMut(CallbackMsg<'a>) -> CallbackReturn> = Box::new(callback);
unsafe {
yara_sys::yr_scanner_set_callback(
scanner,
Some(scan_callback),
&p_callback as *const Box<dyn FnMut(CallbackMsg<'a>) -> CallbackReturn> as *mut _,
);
yara_sys::yr_scanner_scan_fd(scanner, fd)
};
let _: Box<Box<dyn FnMut(CallbackMsg<'a>) -> CallbackReturn>> =
unsafe { Box::from_raw(user_data as *mut _) };
result
}
}

#[cfg(windows)]
Expand All @@ -225,16 +216,15 @@ pub fn scanner_scan_raw<'a>(
callback: impl FnMut(CallbackMsg<'a>) -> CallbackReturn,
) -> i32 {
let handle = file.as_raw_handle();
let p_callback: Box<Box<dyn FnMut(CallbackMsg<'a>) -> CallbackReturn>> =
Box::new(Box::new(callback));
let user_data = Box::into_raw(p_callback) as *mut c_void;
let result = unsafe {
yara_sys::yr_scanner_set_callback(scanner, Some(scan_callback), user_data);
let p_callback: Box<dyn FnMut(CallbackMsg<'a>) -> CallbackReturn> = Box::new(callback);
unsafe {
yara_sys::yr_scanner_set_callback(
scanner,
Some(scan_callback),
&p_callback as *const Box<dyn FnMut(CallbackMsg<'a>) -> CallbackReturn> as *mut _,
);
yara_sys::yr_scanner_scan_fd(scanner, handle)
};
let _: Box<Box<dyn FnMut(CallbackMsg<'a>) -> CallbackReturn>> =
unsafe { Box::from_raw(user_data as *mut _) };
result
}
}

/// Attach a process, pause it, and scan its memory.
Expand All @@ -245,21 +235,17 @@ pub fn rules_scan_proc<'a>(
flags: i32,
callback: impl FnMut(CallbackMsg<'a>) -> CallbackReturn,
) -> Result<(), YaraError> {
let p_callback: Box<Box<dyn FnMut(CallbackMsg<'a>) -> CallbackReturn>> =
Box::new(Box::new(callback));
let user_data = Box::into_raw(p_callback) as *mut c_void;
let p_callback: Box<dyn FnMut(CallbackMsg<'a>) -> CallbackReturn> = Box::new(callback);
let result = unsafe {
yara_sys::yr_rules_scan_proc(
rules,
pid as i32,
flags,
Some(scan_callback),
user_data,
&p_callback as *const Box<dyn FnMut(CallbackMsg<'a>) -> CallbackReturn> as *mut _,
timeout,
)
};
let _: Box<Box<dyn FnMut(CallbackMsg<'a>) -> CallbackReturn>> =
unsafe { Box::from_raw(user_data as *mut _) };

yara_sys::Error::from_code(result)
.map_err(|e| e.into())
Expand Down
90 changes: 76 additions & 14 deletions src/scanner.rs
View file
@@ -1,9 +1,11 @@
use std::fs::File;
use std::marker::PhantomData;
use std::path::Path;

pub use yara_sys::scan_flags::*;

use crate::{compiler::CompilerVariableValue, errors::*, internals, rules::Rules, Rule};
use crate::{compiler::CompilerVariableValue, errors::*, internals, Rule, rules::Rules};
pub use crate::internals::{CallbackMsg, CallbackReturn};

/// A wrapper around compiled [Rules], with its own set of external variables, flags and timeout.
///
Expand Down Expand Up @@ -131,7 +133,26 @@ impl<'rules> Scanner<'rules> {
}
internals::CallbackReturn::Continue
};
internals::scanner_scan_mem(self.inner, mem, callback).map(|_| results)
self.scan_mem_callback(mem, callback).map(|_| results)
}

/// Scan memory with custom callback
///
/// Returns
///
/// * `mem` - Slice to scan
/// * `callback` - YARA callback mor read [here](https://yara.readthedocs.io/en/stable/capi.html#scanning-data)
///
/// # Ownership
///
/// This funciton takes the Scanner as `&mut` because it modifies the
/// `scanner->callback` and `scanner->user_data`, which are not behind a Mutex.
pub fn scan_mem_callback<'r>(
&mut self,
mem: &[u8],
callback: impl FnMut(CallbackMsg<'r>) -> CallbackReturn,
) -> Result<(), YaraError> {
internals::scanner_scan_mem(self.inner, mem, callback)
}

/// Scan a file.
Expand All @@ -142,20 +163,38 @@ impl<'rules> Scanner<'rules> {
///
/// This function takes the Scanner as `&mut` because it modifies the
/// `scanner->callback` and `scanner->user_data`, which are not behind a Mutex.
pub fn scan_file<'r, P: AsRef<Path>>(&self, path: P) -> Result<Vec<Rule<'r>>, Error> {
pub fn scan_file<'r, P: AsRef<Path>>(&mut self, path: P) -> Result<Vec<Rule<'r>>, Error> {
let mut results: Vec<Rule> = Vec::new();
let callback = |message| {
if let internals::CallbackMsg::RuleMatching(rule) = message {
results.push(rule)
}
internals::CallbackReturn::Continue
};

self.scan_file_callback(path, callback).map(|_| results)
}

/// Scan file with custom callback
///
/// Returns
///
/// * `path` - Path to file
/// * `callback` - YARA callback mor read [here](https://yara.readthedocs.io/en/stable/capi.html#scanning-data)
///
/// # Ownership
///
/// This function takes the Scanner as `&mut` because it modifies the
/// `scanner->callback` and `scanner->user_data`, which are not behind a Mutex.
pub fn scan_file_callback<'r, P: AsRef<Path>>(
&mut self,
path: P,
callback: impl FnMut(CallbackMsg<'r>) -> CallbackReturn,
) -> Result<(), Error> {
File::open(path)
.map_err(|e| IoError::new(e, IoErrorKind::OpenScanFile).into())
.and_then(|file| {
let mut results: Vec<Rule> = Vec::new();
let callback = |message| {
if let internals::CallbackMsg::RuleMatching(rule) = message {
results.push(rule)
}
internals::CallbackReturn::Continue
};
internals::scanner_scan_file(self.inner, &file, callback)
.map_err(|e| e.into())
.map(|_| results)
internals::scanner_scan_file(self.inner, &file, callback).map_err(|e| e.into())
})
}

Expand All @@ -179,7 +218,30 @@ impl<'rules> Scanner<'rules> {
}
internals::CallbackReturn::Continue
};
internals::scanner_scan_proc(self.inner, pid, callback).map(|_| results)
self.scan_process_callback(pid, callback).map(|_| results)
}

/// Attach a process, pause it, and scan its memory.
///
/// Returns
///
/// * `pid` - Process id
/// * `callback` - YARA callback mor read [here](https://yara.readthedocs.io/en/stable/capi.html#scanning-data)
///
/// # Permissions
///
/// You need to be able to attach to process `pid`.
///
/// # Ownership
///
/// This function takes the Scanner as `&mut` because it modifies the
/// `scanner->callback` and `scanner->user_data`, which are not behind a Mutex.
pub fn scan_process_callback<'r>(
&mut self,
pid: u32,
callback: impl FnMut(CallbackMsg<'r>) -> CallbackReturn,
) -> Result<(), YaraError> {
internals::scanner_scan_proc(self.inner, pid, callback)
}

/// Set the maximum number of seconds that the scanner will spend in any call
Expand Down

0 comments on commit 1b52d34

Please sign in to comment.