diff --git a/core/src/workspace/config.rs b/core/src/workspace/config.rs index 8107f19..fa604b0 100644 --- a/core/src/workspace/config.rs +++ b/core/src/workspace/config.rs @@ -6,32 +6,15 @@ use crate::VA; use lancelot_flirt::{FlirtSignature, FlirtSignatureSet}; pub trait Configuration: Send { + /// provide the FLIRT signatures to be used to recognize known code. fn get_sigs(&self) -> Result; + + /// provide the addresses known to be functions. fn get_function_hints(&self) -> Result>; + fn clone(&self) -> Box; } -// dummy configuration with only empty values. -pub struct EmptyConfiguration {} - -impl Configuration for EmptyConfiguration { - fn get_sigs(&self) -> Result { - Ok(FlirtSignatureSet::with_signatures(vec![])) - } - - fn get_function_hints(&self) -> Result> { - Ok(vec![]) - } - - fn clone(&self) -> Box { - Box::new(EmptyConfiguration {}) - } -} - -pub fn empty() -> Box { - Box::new(EmptyConfiguration {}) -} - /// Directory that contains: /// - sigs/ FLIRT signatures, ending with .sig, .pat, .sig.gz, .pat.gz pub struct FileSystemConfiguration { @@ -134,3 +117,7 @@ impl Configuration for DynamicConfiguration { }) } } + +pub fn empty() -> Box { + Box::new(DynamicConfiguration::default()) +} diff --git a/pylancelot/src/lib.rs b/pylancelot/src/lib.rs index cf4561f..542c7b1 100644 --- a/pylancelot/src/lib.rs +++ b/pylancelot/src/lib.rs @@ -6,10 +6,15 @@ use ::lancelot::{ module::ModuleError, pagemap::PageMapError, util::UtilError, - workspace::{export::binexport2::export_workspace_to_binexport2, WorkspaceError}, + workspace::{ + config::{Configuration, DynamicConfiguration}, + export::binexport2::export_workspace_to_binexport2, + WorkspaceError, + }, }; use anyhow::Error; use pyo3::{prelude::*, types::*, wrap_pyfunction}; +use std::path::PathBuf; /// ValueError -> "you're doing something wrong" fn to_value_error(e: anyhow::Error) -> PyErr { @@ -61,17 +66,32 @@ fn to_py_err(e: Error) -> PyErr { /// /// Args: /// buf (bytes): the raw bytes of a supported file (e.g., PE or COFF) -/// executable_id (Optional[str]): name of the file +/// executable_id (Optional[str]): name of the file, if known +/// sig_paths (Optional[list[str]]): paths to FLIRT signature files +/// function_hints (Optional[list[int]]): known function virtual addresses /// /// Returns: bytes #[pyfunction] -#[pyo3(signature = (buf, executable_id=None))] +#[pyo3(signature = (buf, executable_id=None, sig_paths=None, function_hints=None))] pub fn binexport2_from_bytes( py: Python, buf: &Bound<'_, PyBytes>, executable_id: Option, + sig_paths: Option>, + function_hints: Option>, ) -> PyResult> { - let config = ::lancelot::workspace::config::empty(); + let mut config: DynamicConfiguration = Default::default(); + if let Some(sig_paths) = sig_paths { + let sig_paths: Vec<_> = sig_paths.iter().map(PathBuf::from).collect(); + config = config.with_sig_paths(&sig_paths); + } + + if let Some(function_hints) = function_hints { + config = config.with_function_hints(&function_hints); + } + + let config = config.clone(); + let ws = ::lancelot::workspace::workspace_from_bytes(config, buf.as_bytes()).map_err(to_py_err)?; let hash = sha256::digest(buf.as_bytes()); export_workspace_to_binexport2(&*ws, hash, executable_id) diff --git a/pylancelot/tests/test_pylancelot.py b/pylancelot/tests/test_pylancelot.py index 0f7e1d9..79202a9 100644 --- a/pylancelot/tests/test_pylancelot.py +++ b/pylancelot/tests/test_pylancelot.py @@ -28,3 +28,9 @@ def test_load_pe(k32): def test_load_coff(altsvc): lancelot.binexport2_from_bytes(altsvc) + +def test_hint_function(k32): + # 7dd70e00 int32_t* __stdcall _GetStartupInfoA@4(int32_t* arg1) + # 7dd70e00 8bff mov edi, edi + # 7dd70e02 55 push ebp {__saved_ebp} + lancelot.binexport2_from_bytes(k32, function_hints=[0x7dd70e02])