Several improvements to python module

boreal-py/src/lib.rs

@@ -11,7 +11,6 @@ use pyo3::{create_exception, ffi, intern};
 use ::boreal::compiler;
 
 // TODO: all clone impls should be efficient...
-// TODO: should all pyclasses have names and be exposed in the module?
 // TODO: check GIL handling in all functions (especially match)
 
 mod rule;
@@ -24,12 +23,29 @@ create_exception!(boreal, AddRuleError, PyException, "error when adding rules");
 
 #[pymodule]
 fn boreal(m: &Bound<'_, PyModule>) -> PyResult<()> {
+    let py = m.py();
+
     m.add_function(wrap_pyfunction!(compile, m)?)?;
-    m.add_function(wrap_pyfunction!(available_modules, m)?)?;
 
-    m.add("AddRuleError", m.py().get_type::<AddRuleError>())?;
-    m.add("ScanError", m.py().get_type::<scanner::ScanError>())?;
-    m.add("TimeoutError", m.py().get_type::<scanner::TimeoutError>())?;
+    m.add("modules", get_available_modules(py))?;
+    m.add("__version__", env!("CARGO_PKG_VERSION"))?;
+
+    m.add("AddRuleError", py.get_type::<AddRuleError>())?;
+    m.add("ScanError", py.get_type::<scanner::ScanError>())?;
+    m.add("TimeoutError", py.get_type::<scanner::TimeoutError>())?;
+
+    m.add("Rule", py.get_type::<rule::Rule>())?;
+    m.add("Match", py.get_type::<rule_match::Match>())?;
+    m.add("Scanner", py.get_type::<scanner::Scanner>())?;
+    m.add("RulesIter", py.get_type::<scanner::RulesIter>())?;
+    m.add(
+        "StringMatchInstance",
+        py.get_type::<string_match_instance::StringMatchInstance>(),
+    )?;
+    m.add(
+        "StringMatches",
+        py.get_type::<string_matches::StringMatches>(),
+    )?;
 
     Ok(())
 }
@@ -150,8 +166,7 @@ fn compile(
     Ok(scanner::Scanner::new(compiler.into_scanner(), warnings))
 }
 
-#[pyfunction]
-fn available_modules(py: Python<'_>) -> Vec<Bound<'_, PyString>> {
+fn get_available_modules(py: Python<'_>) -> Vec<Bound<'_, PyString>> {
     build_compiler()
         .available_modules()
         .map(|s| PyString::new(py, s))

boreal-py/src/rule.rs

@@ -5,7 +5,7 @@ use ::boreal::scanner;
 use ::boreal::{Metadata, MetadataValue};
 
 /// A matching rule
-#[pyclass(frozen)]
+#[pyclass(frozen, module = "boreal")]
 pub struct Rule {
     /// Name of the rule
     #[pyo3(get)]

boreal-py/src/rule_match.rs

@@ -10,7 +10,7 @@ use crate::rule::convert_metadata;
 use crate::string_matches::StringMatches;
 
 /// A matching rule
-#[pyclass(frozen)]
+#[pyclass(frozen, module = "boreal")]
 pub struct Match {
     /// Name of the matching rule
     #[pyo3(get)]

boreal-py/src/scanner.rs

@@ -14,22 +14,18 @@ use crate::rule_match::Match;
 create_exception!(boreal, ScanError, PyException, "error when scanning");
 create_exception!(boreal, TimeoutError, PyException, "scan timed out");
 
-#[pyclass]
+#[pyclass(frozen, module = "boreal")]
 pub struct Scanner {
     scanner: scanner::Scanner,
+
+    /// List of warnings generated when compiling rules.
     #[pyo3(get)]
     warnings: Vec<String>,
-
-    rules_iter: Option<std::vec::IntoIter<crate::rule::Rule>>,
 }
 
 impl Scanner {
     pub fn new(scanner: scanner::Scanner, warnings: Vec<String>) -> Self {
-        Self {
-            scanner,
-            warnings,
-            rules_iter: None,
-        }
+        Self { scanner, warnings }
     }
 }
 
@@ -140,22 +136,34 @@ impl Scanner {
         }
     }
 
-    fn __iter__(mut slf: PyRefMut<Self>) -> PyResult<PyRefMut<Self>> {
+    fn __iter__(&self, py: Python<'_>) -> PyResult<RulesIter> {
         // Unfortunately, we cannot return an object with a lifetime, so
         // we need to collect all rules into a vec of owned elements before
         // generating an iterator...
-        slf.rules_iter = Some(
-            slf.scanner
+        Ok(RulesIter {
+            rules_iter: self
+                .scanner
                 .rules()
-                .map(|rule| crate::rule::Rule::new(slf.py(), &slf.scanner, &rule))
+                .map(|rule| crate::rule::Rule::new(py, &self.scanner, &rule))
                 .collect::<Result<Vec<_>, _>>()?
                 .into_iter(),
-        );
-        Ok(slf)
+        })
+    }
+}
+
+#[pyclass(module = "boreal")]
+pub struct RulesIter {
+    rules_iter: std::vec::IntoIter<crate::rule::Rule>,
+}
+
+#[pymethods]
+impl RulesIter {
+    fn __iter__(slf: PyRef<Self>) -> PyRef<Self> {
+        slf
     }
 
     fn __next__(mut slf: PyRefMut<Self>) -> Option<crate::rule::Rule> {
-        slf.rules_iter.as_mut().and_then(Iterator::next)
+        slf.rules_iter.next()
     }
 }
 

boreal-py/src/string_match_instance.rs

@@ -5,7 +5,7 @@ use pyo3::prelude::*;
 use ::boreal::scanner;
 
 /// Match instance of a YARA string
-#[pyclass(frozen)]
+#[pyclass(frozen, module = "boreal")]
 #[derive(Clone)]
 pub struct StringMatchInstance {
     /// Offset of the match.

boreal-py/src/string_matches.rs

@@ -7,7 +7,7 @@ use ::boreal::scanner;
 use crate::string_match_instance::StringMatchInstance;
 
 /// List of match instances of a YARA string
-#[pyclass(frozen)]
+#[pyclass(frozen, module = "boreal")]
 #[derive(Clone)]
 pub struct StringMatches {
     /// Name of the matching string.

boreal-py/tests/test_api.py

@@ -0,0 +1,17 @@
+import pytest
+from .utils import MODULES
+
+
+@pytest.mark.parametrize('module,is_yara', MODULES)
+def test_modules(module, is_yara):
+    modules = module.modules
+
+    assert type(modules) is list
+    assert 'pe' in modules
+    assert 'time' in modules
+    assert 'console' in modules
+
+
+@pytest.mark.parametrize('module,is_yara', MODULES)
+def test_version(module, is_yara):
+    assert type(module.__version__) is str

boreal-py/tests/test_compile.py

@@ -2,12 +2,7 @@
 import pytest
 import tempfile
 import yara
-
-
-MODULES = [
-    (boreal, False),
-    (yara, True),
-]
+from .utils import MODULES
 
 
 def compile_exc_type(is_yara):

boreal-py/tests/test_scanner.py

@@ -6,12 +6,7 @@
 import subprocess
 import tempfile
 import yara
-
-
-MODULES = [
-    (boreal, False),
-    (yara, True),
-]
+from .utils import MODULES
 
 
 def get_rules(module):
@@ -250,10 +245,10 @@ def test_match_console_log(module, is_yara, capsys):
     assert captured.err == ""
 
 
-# TODO: find a way to compile yara python with cuckoo?
-def test_match_modules_data():
+@pytest.mark.parametrize('module,is_yara', MODULES)
+def test_match_modules_data(module, is_yara):
     # Test only works if the cuckoo module is present
-    if 'cuckoo' not in boreal.available_modules():
+    if 'cuckoo' not in module.modules:
         return
 
     rules = boreal.compile(source="""
@@ -277,7 +272,7 @@ def test_match_modules_data_errors():
     with pytest.raises(TypeError):
         rules.match(data="", modules_data={ 'unknown': 1 })
 
-    if 'cuckoo' in boreal.available_modules():
+    if 'cuckoo' in boreal.modules:
         with pytest.raises(TypeError):
             rules.match(data="", modules_data={ 'cuckoo': 1 })
         with pytest.raises(TypeError):

boreal-py/tests/test_types.py

@@ -1,12 +1,5 @@
-import boreal
 import pytest
-import yara
-
-
-MODULES = [
-    (boreal, False),
-    (yara, True),
-]
+from .utils import MODULES
 
 
 @pytest.mark.parametrize("module,is_yara", MODULES)

boreal-py/tests/utils.py

@@ -0,0 +1,8 @@
+import boreal
+import yara
+
+
+MODULES = [
+    (boreal, False),
+    (yara, True),
+]