feat: load yara lazily to avoid action dispatcher error (#1162)

- Implement lazy loading of yara module with proper type hints
- Add error messages when yara is not available
- Add helper function _check_yara_available() for consistent error handling
- make helper functions private and remove checks
 - add test for yara import error handling to appreciate codecov

Author: Pouyanpi

nemoguardrails/library/injection_detection/actions.py

@@ -34,7 +34,11 @@
 from pathlib import Path
 from typing import Tuple, Union
 
-import yara
+yara = None
+try:
+    import yara
+except ImportError:
+    pass
 
 from nemoguardrails import RailsConfig
 from nemoguardrails.actions import action
@@ -45,6 +49,14 @@
 log = logging.getLogger(__name__)
 
 
+def _check_yara_available():
+    if yara is None:
+        raise ImportError(
+            "The yara module is required for injection detection. "
+            "Please install it using: pip install yara-python"
+        )
+
+
 def _validate_unpack_config(config: RailsConfig) -> Tuple[str, Path, Tuple[str]]:
     """
     Validates and unpacks the injection detection configuration.
@@ -117,7 +129,7 @@ def _validate_unpack_config(config: RailsConfig) -> Tuple[str, Path, Tuple[str]]
 
 
 @lru_cache()
-def load_rules(yara_path: Path, rule_names: Tuple) -> Union[yara.Rules, None]:
+def _load_rules(yara_path: Path, rule_names: Tuple) -> Union["yara.Rules", None]:
     """
     Loads and compiles YARA rules from the specified path and rule names.
 
@@ -126,12 +138,14 @@ def load_rules(yara_path: Path, rule_names: Tuple) -> Union[yara.Rules, None]:
         rule_names (Tuple): A tuple of YARA rule names to load.
 
     Returns:
-        Union[yara.Rules, None]: The compiled YARA rules object if successful,
+        Union['yara.Rules', None]: The compiled YARA rules object if successful,
         or None if no rule names are provided.
 
     Raises:
         yara.SyntaxError: If there is a syntax error in the YARA rules.
+        ImportError: If the yara module is not installed.
     """
+
     if len(rule_names) == 0:
         log.warning(
             "Injection config was provided but no modules were specified. Returning None."
@@ -150,7 +164,7 @@ def load_rules(yara_path: Path, rule_names: Tuple) -> Union[yara.Rules, None]:
     return rules
 
 
-def omit_injection(text: str, matches: list[yara.Match]) -> str:
+def _omit_injection(text: str, matches: list["yara.Match"]) -> str:
     """
     Attempts to strip the offending injection attempts from the provided text.
 
@@ -160,11 +174,15 @@ def omit_injection(text: str, matches: list[yara.Match]) -> str:
 
     Args:
         text (str): The text to check for command injection.
-        matches (list[yara.Match]): A list of YARA rule matches.
+        matches (list['yara.Match']): A list of YARA rule matches.
 
     Returns:
         str: The text with the detected injections stripped out.
+
+    Raises:
+        ImportError: If the yara module is not installed.
     """
+
     # Copy the text to a placeholder variable
     modified_text = text
     for match in matches:
@@ -180,7 +198,7 @@ def omit_injection(text: str, matches: list[yara.Match]) -> str:
     return modified_text
 
 
-def sanitize_injection(text: str, matches: list[yara.Match]) -> str:
+def _sanitize_injection(text: str, matches: list["yara.Match"]) -> str:
     """
     Attempts to sanitize the offending injection attempts in the provided text.
     This is done by 'de-fanging' the offending content, transforming it into a state that will not execute
@@ -193,20 +211,22 @@ def sanitize_injection(text: str, matches: list[yara.Match]) -> str:
 
     Args:
         text (str): The text to check for command injection.
-        matches (list[yara.Match]): A list of YARA rule matches.
+        matches (list['yara.Match']): A list of YARA rule matches.
 
     Returns:
         str: The text with the detected injections sanitized.
 
     Raises:
         NotImplementedError: If the sanitization logic is not implemented.
+        ImportError: If the yara module is not installed.
     """
+
     raise NotImplementedError(
         "Injection sanitization is not yet implemented. Please use 'reject' or 'omit'"
     )
 
 
-def reject_injection(text: str, rules: yara.Rules) -> Tuple[bool, str]:
+def _reject_injection(text: str, rules: "yara.Rules") -> Tuple[bool, str]:
     """
     Detects whether the provided text contains potential injection attempts.
 
@@ -215,15 +235,17 @@ def reject_injection(text: str, rules: yara.Rules) -> Tuple[bool, str]:
 
     Args:
         text (str): The text to check for command injection.
-        rules (yara.Rules): The loaded YARA rules.
+        rules ('yara.Rules'): The loaded YARA rules.
 
     Returns:
         bool: True if attempted exploitation is detected, False otherwise.
         str: list of matches as a string
 
     Raises:
         ValueError: If the `action` parameter in the configuration is invalid.
+        ImportError: If the yara module is not installed.
     """
+
     if rules is None:
         log.warning(
             "reject_injection guardrail was invoked but no rules were specified in the InjectionDetection config."
@@ -258,10 +280,12 @@ async def injection_detection(text: str, config: RailsConfig) -> str:
         ValueError: If the `action` parameter in the configuration is invalid.
         NotImplementedError: If an unsupported action is encountered.
     """
+    _check_yara_available()
+
     action_option, yara_path, rule_names = _validate_unpack_config(config)
-    rules = load_rules(yara_path, rule_names)
+    rules = _load_rules(yara_path, rule_names)
     if action_option == "reject":
-        verdict, detections = reject_injection(text, rules)
+        verdict, detections = _reject_injection(text, rules)
         if verdict:
             return f"I'm sorry, the desired output triggered rule(s) designed to mitigate exploitation of {detections}."
         else:
@@ -276,9 +300,9 @@ async def injection_detection(text: str, config: RailsConfig) -> str:
         matches_string = ", ".join([match_name.rule for match_name in matches])
         log.info(f"Input matched on rule {matches_string}.")
         if action_option == "omit":
-            return omit_injection(text, matches)
+            return _omit_injection(text, matches)
         elif action_option == "sanitize":
-            return sanitize_injection(text, matches)
+            return _sanitize_injection(text, matches)
         else:
             # We should never ever hit this since we inspect the action option above, but putting an error here anyway.
             raise NotImplementedError(

tests/test_injection_detection.py

@@ -39,11 +39,12 @@
 from nemoguardrails.actions import action
 from nemoguardrails.actions.actions import ActionResult
 from nemoguardrails.library.injection_detection.actions import (
+    _check_yara_available,
+    _load_rules,
+    _omit_injection,
+    _reject_injection,
     _validate_unpack_config,
     injection_detection,
-    load_rules,
-    omit_injection,
-    reject_injection,
 )
 from tests.utils import TestChat
 
@@ -101,14 +102,14 @@ def test_load_custom_rules():
             """,
     )
     _action_option, yara_path, rule_names = _validate_unpack_config(config)
-    rules = load_rules(yara_path, rule_names)
+    rules = _load_rules(yara_path, rule_names)
     assert isinstance(rules, yara.Rules)
 
 
 def test_load_all_rules():
     config = RailsConfig.from_path(os.path.join(CONFIGS_FOLDER, "injection_detection"))
     _action_option, yara_path, rule_names = _validate_unpack_config(config)
-    rules = load_rules(yara_path, rule_names)
+    rules = _load_rules(yara_path, rule_names)
     assert isinstance(rules, yara.Rules)
 
 
@@ -204,7 +205,7 @@ def test_empty_injection_rules():
             """
     )
     _action_option, yara_path, rule_names = _validate_unpack_config(config)
-    rules = load_rules(yara_path, rule_names)
+    rules = _load_rules(yara_path, rule_names)
     assert rules is None
 
 
@@ -220,7 +221,7 @@ async def test_omit_injection_action():
         create_mock_yara_match("-- comment", "sqli"),
     ]
 
-    result = omit_injection(text=text, matches=mock_matches)
+    result = _omit_injection(text=text, matches=mock_matches)
 
     # all sql injection should be removed
     # NOTE: following rule does not get removed using sqli.yara
@@ -243,11 +244,11 @@ async def test_reject_injection_with_mismatched_action():
 
     # pathcing the load_rules function to return our mock rules
     with patch(
-        "nemoguardrails.library.injection_detection.actions.load_rules",
+        "nemoguardrails.library.injection_detection.actions._load_rules",
         return_value=mock_rules,
     ):
         sql_injection = "' OR 1 = 1"
-        result, _ = reject_injection(sql_injection, mock_rules)
+        result, _ = _reject_injection(sql_injection, mock_rules)
         assert result is True
 
 
@@ -263,11 +264,11 @@ async def test_multiple_injection_types():
     mock_rules = create_mock_rules(mock_matches)
 
     with patch(
-        "nemoguardrails.library.injection_detection.actions.load_rules",
+        "nemoguardrails.library.injection_detection.actions._load_rules",
         return_value=mock_rules,
     ):
         multi_injection = "' OR 1 = 1 <script>alert('xss')</script>"
-        result, _ = reject_injection(multi_injection, mock_rules)
+        result, _ = _reject_injection(multi_injection, mock_rules)
         assert result is True
 
 
@@ -279,21 +280,21 @@ async def test_edge_cases():
     mock_rules = create_mock_rules([])
 
     with patch(
-        "nemoguardrails.library.injection_detection.actions.load_rules",
+        "nemoguardrails.library.injection_detection.actions._load_rules",
         return_value=mock_rules,
     ):
         # Test with empty string
-        result, _ = reject_injection("", mock_rules)
+        result, _ = _reject_injection("", mock_rules)
         assert result is False
 
         # no issue with very long str
         long_string = "a" * 10000
-        result, _ = reject_injection(long_string, mock_rules)
+        result, _ = _reject_injection(long_string, mock_rules)
         assert result is False
 
         # no issue with special chars
         special_chars = "!@#$%^&*()_+-=[]{}|;:,.<>?"
-        result, _ = reject_injection(special_chars, mock_rules)
+        result, _ = _reject_injection(special_chars, mock_rules)
         assert result is False
 
 
@@ -495,3 +496,18 @@ async def test_sanitize_action_not_implemented():
 
                     """
         )
+
+
+def test_yara_import_error():
+    """Test that appropriate error is raised when yara module is not available."""
+
+    with patch("nemoguardrails.library.injection_detection.actions.yara", None):
+        with pytest.raises(ImportError) as exc_info:
+            _check_yara_available()
+        assert str(exc_info.value) == (
+            "The yara module is required for injection detection. "
+            "Please install it using: pip install yara-python"
+        )
+
+    with patch("nemoguardrails.library.injection_detection.actions.yara", yara):
+        _check_yara_available()