feat: expose linter API of compiler to python bindings. (VirusTotal#612)

wxsBSD · plusvic · web-flow · commit cd80efa2ddb4 · 2026-03-31T12:33:05.000+02:00
This commit exposes the linters to the python API via the Compiler class. With
this you can now do things like:

```
import yara_x

rule = '''
rule test : a b c {
  meta:
    author = "foo"
  strings:
    $a = "bar"
  condition:
    $a
}'''

c = yara_x.Compiler()
c.allowed_metadata('author', yara_x.MetaType.STRING, regexp='bar')
c.allowed_tags(['a', 'b'], error = True)
```

---------

Co-authored-by: Victor M. Alvarez &lt;vmalvarez@virustotal.com&gt;
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/py/Cargo.toml b/py/Cargo.toml
@@ -60,13 +60,15 @@ pyo3 = { version = "0.28.2", features = [
     "abi3-py38",
     "extension-module",
 ] }
+regex = { workspace = true }
 serde_json = { workspace = true }
 strum = { workspace = true }
 strum_macros = { workspace = true }
 
 yara-x = { workspace = true }
 yara-x-proto-json = { workspace = true }
 yara-x-fmt = { workspace = true }
+yara-x-parser = { workspace = true }
 
 [build-dependencies]
 pyo3-build-config = "0.28.2"
diff --git a/py/src/lib.rs b/py/src/lib.rs
@@ -40,6 +40,7 @@ use strum_macros::{Display, EnumString};
 
 use ::yara_x as yrx;
 use yara_x_fmt::Indentation;
+use yara_x_parser::ast::MetaValue;
 
 fn dict_to_json(dict: Bound<PyAny>) -> PyResult<serde_json::Value> {
     static JSON_DUMPS: PyOnceLock<Py<PyAny>> = PyOnceLock::new();
@@ -72,6 +73,36 @@ enum SupportedModules {
     Dex,
 }
 
+// These are copies from the checker in the CLI, but exposing them in the API
+// for use here seems wrong. Maybe move them to a better place or just keep our
+// own copies here?
+fn is_sha256(s: &str) -> bool {
+    s.len() == 64 && s.chars().all(|c| c.is_ascii_hexdigit())
+}
+
+fn is_sha1(s: &str) -> bool {
+    s.len() == 40 && s.chars().all(|c| c.is_ascii_hexdigit())
+}
+
+fn is_md5(s: &str) -> bool {
+    s.len() == 32 && s.chars().all(|c| c.is_ascii_hexdigit())
+}
+
+/// Supported metadata types used to add linters to the compiler.
+#[pyclass(from_py_object)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+#[allow(clippy::upper_case_acronyms)]
+enum MetaType {
+    STRING,
+    INTEGER,
+    FLOAT,
+    BOOL,
+    SHA256,
+    SHA1,
+    MD5,
+    HASH,
+}
+
 /// Formats YARA rules.
 #[pyclass(unsendable)]
 struct Formatter {
@@ -420,10 +451,17 @@ impl Compiler {
     /// will return an InvalidRuleName warning.
     ///
     /// If the regexp does not compile a ValueError is returned.
-    #[pyo3(signature = (regexp))]
-    fn rule_name_regexp(&mut self, regexp: &str) -> PyResult<()> {
-        let linter = yrx::linters::rule_name(regexp)
-            .map_err(|err| PyValueError::new_err(err.to_string()))?;
+    #[pyo3(signature = (regexp, error = false))]
+    fn allowed_rule_name(
+        &mut self,
+        regexp: &str,
+        error: bool,
+    ) -> PyResult<()> {
+        let mut linter = match yrx::linters::rule_name(regexp) {
+            Ok(linter) => linter,
+            Err(err) => return Err(PyValueError::new_err(err.to_string())),
+        };
+        linter = linter.error(error);
         self.inner.add_linter(linter);
         Ok(())
     }
@@ -616,6 +654,125 @@ impl Compiler {
             .map_err(|err| PyValueError::new_err(err.to_string()))?;
         json_loads.call((warnings_json,), None)
     }
+
+    #[pyo3(signature = (tags, error = false))]
+    fn allowed_tags(
+        &mut self,
+        tags: Vec<String>,
+        error: bool,
+    ) -> PyResult<()> {
+        self.inner.add_linter(yrx::linters::tags_allowed(tags).error(error));
+        Ok(())
+    }
+
+    #[pyo3(signature = (regexp, error = false))]
+    fn allowed_tags_regex(
+        &mut self,
+        regexp: String,
+        error: bool,
+    ) -> PyResult<()> {
+        let mut linter = match yrx::linters::tag_regex(regexp) {
+            Ok(linter) => linter,
+            Err(err) => return Err(PyValueError::new_err(err.to_string())),
+        };
+        linter = linter.error(error);
+        self.inner.add_linter(linter);
+        Ok(())
+    }
+
+    #[pyo3(signature = (
+            identifier,
+            value_type,
+            required = false,
+            error = false,
+            regexp = None
+        ))]
+    fn allowed_metadata(
+        &mut self,
+        identifier: &str,
+        value_type: MetaType,
+        required: bool,
+        error: bool,
+        regexp: Option<String>,
+    ) -> PyResult<()> {
+        let mut linter =
+            yrx::linters::metadata(identifier).required(required).error(error);
+        match value_type {
+            MetaType::STRING => {
+                let message = if let Some(regexp) = regexp.clone() {
+                    let _ = regex::bytes::Regex::new(regexp.as_str())
+                        .map_err(|err| PyValueError::new_err(err.to_string()));
+                    format!(
+                        "`{identifier}` must be a string that matches `/{regexp}/`"
+                    )
+                } else {
+                    format!("`{identifier}` must be a string")
+                };
+                linter = linter.validator(
+                    move |meta| match (&meta.value, &regexp) {
+                        (MetaValue::String((s, _)), Some(regexp)) => {
+                            regex::Regex::new(regexp.as_str())
+                                .unwrap()
+                                .is_match(s)
+                        }
+                        (MetaValue::Bytes((s, _)), Some(regexp)) => {
+                            regex::bytes::Regex::new(regexp.as_str())
+                                .unwrap()
+                                .is_match(s)
+                        }
+                        (MetaValue::String(_), None) => true,
+                        (MetaValue::Bytes(_), None) => true,
+                        _ => false,
+                    },
+                    message,
+                );
+            }
+            MetaType::INTEGER => {
+                linter = linter.validator(
+                    |meta| matches!(meta.value, MetaValue::Integer(_)),
+                    format!("`{identifier}` must be an integer"),
+                );
+            }
+            MetaType::FLOAT => {
+                linter = linter.validator(
+                    |meta| matches!(meta.value, MetaValue::Float(_)),
+                    format!("`{identifier}` must be a float"),
+                );
+            }
+            MetaType::BOOL => {
+                linter = linter.validator(
+                    |meta| matches!(meta.value, MetaValue::Bool(_)),
+                    format!("`{identifier}` must be a bool"),
+                );
+            }
+            MetaType::SHA256 => {
+                linter = linter.validator(
+                        |meta| matches!(meta.value, MetaValue::String((s,_)) if is_sha256(s)),
+                        format!("`{identifier}` must be a SHA-256"),
+                    );
+            }
+            MetaType::SHA1 => {
+                linter = linter.validator(
+                        |meta| matches!(meta.value, MetaValue::String((s,_)) if is_sha1(s)),
+                        format!("`{identifier}` must be a SHA-1"),
+                    );
+            }
+            MetaType::MD5 => {
+                linter = linter.validator(
+                        |meta| matches!(meta.value, MetaValue::String((s,_)) if is_md5(s)),
+                        format!("`{identifier}` must be a MD5"),
+                    );
+            }
+            MetaType::HASH => {
+                linter = linter.validator(
+                        |meta| matches!(meta.value, MetaValue::String((s,_)) if is_md5(s) || is_sha1(s) || is_sha256(s)),
+                        format!("`{identifier}` must be a MD5, SHA-1 or SHA-256"),
+                    );
+            }
+        }
+        self.inner.add_linter(linter);
+        Ok(())
+    }
 }
 
 /// Optional information for the scan operation.
@@ -1306,6 +1463,7 @@ fn yara_x(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<Match>()?;
     m.add_class::<Formatter>()?;
     m.add_class::<Module>()?;
+    m.add_class::<MetaType>()?;
     m.gil_used(false)?;
     Ok(())
 }
diff --git a/py/tests/test_api.py b/py/tests/test_api.py
@@ -32,7 +32,7 @@ def test_error_on_slow_pattern():
 def test_invalid_rule_name_regexp():
   compiler = yara_x.Compiler()
   with pytest.raises(ValueError):
-    compiler.rule_name_regexp("(AXS|ERS")
+    compiler.allowed_rule_name("(AXS|ERS")
 
 
 def test_int_globals():
@@ -397,3 +397,60 @@ def test_rules_imports():
 }
 ''')
   assert rules.imports() == ["pe", "elf"]
+
+def test_check_allowed_tags_error():
+  rule = '''
+  rule test: a b c d { condition: 1 + 1 == 2}
+  rule test2: d { condition: 1 + 1 == 2}'''
+  compiler = yara_x.Compiler()
+  compiler.allowed_tags(['a', 'b'], error = True)
+  with pytest.raises(yara_x.CompileError,
+                     match="tag `c` not in allowed list"):
+    compiler.add_source(rule)
+  # The current behavior is stop checking tags on the rule after the first tag
+  # fails, but subsequent rules are also checked.
+  errors = compiler.errors()
+  assert len(errors) == 2
+  assert 'tag `c` not in allowed list' in errors[0]['text']
+  assert 'tag `d` not in allowed list' in errors[1]['text']
+
+def test_check_allowed_tags_warning():
+  compiler = yara_x.Compiler()
+  compiler.allowed_tags(['a', 'b'])
+  compiler.add_source('rule test: a b c d { condition: 1 + 1 == 2}')
+  warnings = compiler.warnings()
+  assert len(warnings) == 2
+  assert 'tag `c` not in allowed list' in warnings[0]['text']
+  assert 'tag `d` not in allowed list' in warnings[1]['text']
+
+def test_check_metadata():
+  compiler = yara_x.Compiler()
+  compiler.allowed_metadata('a', yara_x.MetaType.STRING)
+  compiler.allowed_metadata('b', yara_x.MetaType.STRING, regexp='^bar')
+  compiler.add_source('rule test { meta: a = 1 b = "foo" condition: 1 + 1 == 2}')
+  warnings = compiler.warnings()
+  assert len(warnings) == 2
+  assert '`a` must be a string' in warnings[0]['text']
+  assert '`b` must be a string that matches `/^bar/`' in warnings[1]['text']
+
+def test_check_rule_name_regexp():
+  rule = '''
+  rule test { condition: 1 + 1 == 2}
+  rule test2 { condition: 1 + 1 == 2}'''
+  compiler = yara_x.Compiler()
+  compiler.allowed_rule_name('^foo')
+  compiler.add_source(rule)
+  warnings = compiler.warnings()
+  assert len(warnings) == 2
+  assert 'this rule name does not match regex `^foo`' in warnings[0]['text']
+
+def test_check_rule_name_regexp_error():
+  rule = '''
+  rule test { condition: 1 + 1 == 2}
+  rule test2 { condition: 1 + 1 == 2}'''
+  compiler = yara_x.Compiler()
+  compiler.allowed_rule_name('^foo', error = True)
+  with pytest.raises(yara_x.CompileError,
+                     match=r"this rule name does not match regex `\^foo`"):
+    compiler.add_source(rule)
+  assert len(compiler.errors()) == 2
diff --git a/py/yara_x.pyi b/py/yara_x.pyi
@@ -1,6 +1,7 @@
 import collections
 
-from typing import Any, Dict, BinaryIO, TextIO, Optional, Tuple, final
+from typing import Any, Dict, BinaryIO, TextIO, Optional, Tuple, final, List
+from enum import Enum
 
 class CompileError(Exception):
     r"""
@@ -156,7 +157,7 @@ class Compiler:
         """
         ...
 
-    def rule_name_regexp(self, regexp: str) -> None:
+    def allowed_rule_name(self, regexp: str, error: bool = False) -> None:
         r"""
         Tell the compiler that any rule must match this regular expression or it
         will result in a compiler warning.
@@ -168,6 +169,24 @@ class Compiler:
         """
         ...
 
+    def allowed_tags(self, tags: List[str], error: bool = False) -> None:
+        r"""List the allowed tags for rules."""
+        ...
+
+    def allowed_tags_regex(self, regexp: str, error: bool = False) -> None:
+        r"""A regular expression that must match all tags on rules."""
+        ...
+
+    def allowed_metadata(
+            self,
+            identifier: str,
+            value_type: MetaType,
+            required: bool,
+            regexp: Optional[str],
+            error: bool = False):
+        r"""Define expected type and value for metadata on rules."""
+        ...
+
 @final
 class ScanOptions:
     r"""
@@ -480,3 +499,33 @@ class Module:
     def invoke(self, data: str) -> Any:
         r"""Parse the data and collect module metadata."""
         ...
+
+@final
+class MetaType(Enum):
+    STRING: int
+    INTEGER: int
+    FLOAT: int
+    BOOL: int
+    SHA256: int
+    SHA1: int
+    MD5: int
+    HASH: int
+
+@final
+class CheckResult:
+    r"""Result from the [`Compiler::check`] method after checking source code."""
+    def warning(self) -> bool:
+        r"""True if the result is a warning, false if it is an error."""
+        ...
+
+    def code(self) -> bool:
+        r"""The string representation of the result code."""
+        ...
+
+    def title(self) -> str:
+        r"""The title of the result code."""
+        ...
+
+    def message(self) -> str:
+        r"""A multi-line message containing code, title and full compiler details."""
+        ...
