pylancelot: specify FLIRT paths and known functions via kwargs

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<FlirtSignatureSet>;
+
+    /// provide the addresses known to be functions.
     fn get_function_hints(&self) -> Result<Vec<VA>>;
+
     fn clone(&self) -> Box<dyn Configuration>;
 }
 
-// dummy configuration with only empty values.
-pub struct EmptyConfiguration {}
-
-impl Configuration for EmptyConfiguration {
-    fn get_sigs(&self) -> Result<FlirtSignatureSet> {
-        Ok(FlirtSignatureSet::with_signatures(vec![]))
-    }
-
-    fn get_function_hints(&self) -> Result<Vec<VA>> {
-        Ok(vec![])
-    }
-
-    fn clone(&self) -> Box<dyn Configuration> {
-        Box::new(EmptyConfiguration {})
-    }
-}
-
-pub fn empty() -> Box<dyn Configuration> {
-    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<dyn Configuration> {
+    Box::new(DynamicConfiguration::default())
+}

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<String>,
+    sig_paths: Option<Vec<String>>,
+    function_hints: Option<Vec<u64>>,
 ) -> PyResult<Py<PyBytes>> {
-    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)

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])