feat: implement all options of compile method

Author: vthib

Cargo.lock

@@ -17,6 +17,7 @@ crate-type = ["cdylib"]
 [dependencies]
 boreal = { path = "../boreal", version = "0.9.0" }
 
+libc = "0.2"
 pyo3 = { version = "0.23", features = ["extension-module", "macros"] }
 
 [lints]

boreal-py/Cargo.toml

@@ -1,7 +1,17 @@
 //! Python bindings for the boreal library.
+#![allow(unsafe_code)]
+
+use std::fs::File;
+use std::io::Read;
+use std::os::fd::{FromRawFd, OwnedFd};
+
+use pyo3::create_exception;
+use pyo3::exceptions::{PyException, PyTypeError};
+use pyo3::ffi;
 use pyo3::prelude::*;
+use pyo3::types::PyDict;
 
-use ::boreal::Compiler;
+use ::boreal::compiler;
 
 // TODO: all clone impls should be efficient...
 // TODO: should all pyclasses have names and be exposed in the module?
@@ -11,23 +21,164 @@ mod scanner;
 mod string_match_instance;
 mod string_matches;
 
+create_exception!(boreal, AddRuleError, PyException, "error when adding rules");
+
 #[pymodule]
 fn boreal(m: &Bound<'_, PyModule>) -> PyResult<()> {
-    m.add_function(wrap_pyfunction!(compile, m)?)
+    m.add_function(wrap_pyfunction!(compile, m)?)?;
+    m.add("AddRuleError", m.py().get_type::<AddRuleError>())?;
+    Ok(())
 }
 
+// TODO: add strict_escape?
 #[pyfunction]
-#[pyo3(signature = (filepath=None, source=None))]
-fn compile(filepath: Option<&str>, source: Option<&str>) -> PyResult<scanner::PyScanner> {
-    let mut compiler = Compiler::new();
-    match (filepath, source) {
-        (Some(v), None) => compiler.add_rules_file(v),
-        (None, Some(v)) => compiler.add_rules_str(v),
-        _ => todo!(),
+#[pyo3(signature = (filepath=None, source=None, file=None, filepaths=None, sources=None, externals=None, includes=true, error_on_warning=false))]
+#[allow(clippy::too_many_arguments)]
+fn compile(
+    filepath: Option<&str>,
+    source: Option<&str>,
+    file: Option<&Bound<'_, PyAny>>,
+    filepaths: Option<&Bound<'_, PyDict>>,
+    sources: Option<&Bound<'_, PyDict>>,
+    externals: Option<&Bound<'_, PyDict>>,
+    includes: bool,
+    error_on_warning: bool,
+) -> PyResult<scanner::PyScanner> {
+    let mut compiler = compiler::Compiler::new();
+    compiler.set_params(
+        compiler::CompilerParams::default()
+            .disable_includes(!includes)
+            .fail_on_warnings(error_on_warning),
+    );
+    if let Some(externals) = externals {
+        add_externals(&mut compiler, externals)?;
+    }
+
+    let mut warnings = Vec::new();
+
+    match (filepath, source, file, filepaths, sources) {
+        (Some(filepath), None, None, None, None) => {
+            let res = compiler
+                .add_rules_file(filepath)
+                // TODO: contents
+                .map_err(|err| convert_compiler_error(&err, filepath, ""))
+                .map_err(AddRuleError::new_err)?;
+            warnings = res
+                .warnings()
+                .map(|err| convert_compiler_error(err, filepath, ""))
+                .collect();
+        }
+        (None, Some(source), None, None, None) => {
+            let res = compiler
+                .add_rules_str(source)
+                .map_err(|err| convert_compiler_error(&err, "source", source))
+                .map_err(AddRuleError::new_err)?;
+            warnings = res
+                .warnings()
+                .map(|err| convert_compiler_error(err, "source", source))
+                .collect();
+        }
+        (None, None, Some(file), None, None) => {
+            // Safety:
+            // - `as_ptr` is safe because it does not outlive file
+            // - `PyObject_AsFileDescriptor` is safe to call on any valid PyObject.
+            //   if it fails, -1 is returned.
+            let fd = unsafe { ffi::PyObject_AsFileDescriptor(file.as_ptr()) };
+            if fd == -1 {
+                return Err(PyTypeError::new_err("`file` argument is not a file object"));
+            }
+
+            // Safety: the passed file descriptor is valid and thus can be dupped.
+            let owned_fd = unsafe { libc::dup(fd) };
+            if owned_fd == -1 {
+                return Err(std::io::Error::last_os_error().into());
+            }
+
+            // Safety: the file descriptor is valid and is owned by us.
+            let owned_fd = unsafe { OwnedFd::from_raw_fd(owned_fd) };
+
+            let mut file = File::from(owned_fd);
+            let mut contents = String::new();
+            let _r = file.read_to_string(&mut contents)?;
+            // TODO: this makes the error message not as nice
+            let res = compiler
+                .add_rules_str(&contents)
+                .map_err(|err| convert_compiler_error(&err, "file", &contents))
+                .map_err(AddRuleError::new_err)?;
+            warnings = res
+                .warnings()
+                .map(|err| convert_compiler_error(err, "file", &contents))
+                .collect();
+        }
+        (None, None, None, Some(filepaths), None) => {
+            for (key, value) in filepaths {
+                let namespace: &str = key.extract().map_err(|_| {
+                    PyTypeError::new_err("keys of the `filepaths` argument must be strings")
+                })?;
+                let filepath: &str = value.extract().map_err(|_| {
+                    PyTypeError::new_err("values of the `filepaths` argument must be strings")
+                })?;
+                let res = compiler
+                    .add_rules_file_in_namespace(filepath, namespace)
+                    // TODO: contents
+                    .map_err(|err| convert_compiler_error(&err, filepath, ""))
+                    .map_err(AddRuleError::new_err)?;
+                warnings.extend(
+                    res.warnings()
+                        .map(|err| convert_compiler_error(err, filepath, "")),
+                );
+            }
+        }
+        (None, None, None, None, Some(sources)) => {
+            for (key, value) in sources {
+                let namespace: &str = key.extract().map_err(|_| {
+                    PyTypeError::new_err("keys of the `sources` argument must be strings")
+                })?;
+                let source: &str = value.extract().map_err(|_| {
+                    PyTypeError::new_err("values of the `sources` argument must be strings")
+                })?;
+                let res = compiler
+                    .add_rules_str_in_namespace(source, namespace)
+                    .map_err(|err| convert_compiler_error(&err, namespace, source))
+                    .map_err(AddRuleError::new_err)?;
+                warnings.extend(
+                    res.warnings()
+                        .map(|err| convert_compiler_error(err, namespace, source)),
+                );
+            }
+        }
+        _ => return Err(PyTypeError::new_err("invalid arguments passed")),
     }
-    .unwrap();
 
     Ok(scanner::PyScanner {
         scanner: compiler.into_scanner(),
+        warnings,
     })
 }
+
+fn convert_compiler_error(err: &compiler::AddRuleError, input_name: &str, input: &str) -> String {
+    err.to_short_description(input_name, input)
+}
+
+fn add_externals(compiler: &mut compiler::Compiler, externals: &Bound<'_, PyDict>) -> PyResult<()> {
+    for (key, value) in externals {
+        let name: &str = key.extract()?;
+
+        if let Ok(v) = value.extract::<bool>() {
+            let _r = compiler.define_symbol(name, v);
+        } else if let Ok(v) = value.extract::<i64>() {
+            let _r = compiler.define_symbol(name, v);
+        } else if let Ok(v) = value.extract::<f64>() {
+            let _r = compiler.define_symbol(name, v);
+        } else if let Ok(v) = value.extract::<&str>() {
+            let _r = compiler.define_symbol(name, v);
+        } else if let Ok(v) = value.extract::<&[u8]>() {
+            let _r = compiler.define_symbol(name, v);
+        } else {
+            return Err(PyTypeError::new_err(
+                "invalid type for the external value, must be a boolean, integer, float or string",
+            ));
+        }
+    }
+    Ok(())
+}

boreal-py/src/lib.rs

@@ -9,6 +9,8 @@ use crate::rule_match::Match;
 #[pyclass]
 pub struct PyScanner {
     pub scanner: Scanner,
+    #[pyo3(get)]
+    pub warnings: Vec<String>,
 }
 
 #[pymethods]

boreal-py/src/scanner.rs

@@ -9,54 +9,6 @@
 ]
 
 
-RULE = """
-rule foo: bar baz {
-    meta:
-        s = "a\\nz"
-        b = true
-        v = -11
-    strings:
-        $ = "fgj"
-        $a = /l.{1,2}n/
-    condition:
-        any of them
-}
-"""
-
-@pytest.mark.parametrize("module,is_yara", MODULES)
-def test_overall(module, is_yara):
-    rule = module.compile(source=RULE)
-    matches = rule.match(data='abcdefgjiklmnoprstuvwxyzlmmn')
-    assert len(matches) == 1
-    assert matches[0].rule == 'foo'
-    # FIXME: difference between yara and boreal
-    # assert matches[0].namespace == ''
-    assert matches[0].tags == ['bar', 'baz']
-    assert matches[0].meta == {
-        # XXX yara forces a string type, losing information.
-        's': 'a\nz' if is_yara else b'a\nz',
-        'b': True,
-        'v': -11
-    }
-
-    m = matches[0]
-    assert len(m.strings) == 2
-    assert m.strings[0].identifier == '$'
-    assert len(m.strings[0].instances) == 1
-    assert m.strings[0].instances[0].offset == 5
-    assert m.strings[0].instances[0].matched_length == 3
-    assert m.strings[0].instances[0].matched_data == b'fgj'
-
-    assert m.strings[1].identifier == '$a'
-    assert len(m.strings[1].instances) == 2
-    assert m.strings[1].instances[0].offset == 10
-    assert m.strings[1].instances[0].matched_length == 3
-    assert m.strings[1].instances[0].matched_data == b'lmn'
-    assert m.strings[1].instances[1].offset == 24
-    assert m.strings[1].instances[1].matched_length == 4
-    assert m.strings[1].instances[1].matched_data == b'lmmn'
-
-
 @pytest.mark.parametrize("module,is_yara", MODULES)
 def test_match(module, is_yara):
     """Test all properties related to the Match object"""