fix: yara ingest fails on too many streams. (#36)

fix: yara ingest fails on too many streams.

---------

Co-authored-by: dan-acsc <dan-acsc@users.noreply.github.com>

Author: dan-acsc

azul_plugin_yara/main.py

@@ -6,7 +6,7 @@
 import re
 from hashlib import md5
 from pathlib import Path
-from typing import Any, Optional
+from typing import Annotated, Any, Optional
 
 import yara_x
 from azul_runner import (
@@ -21,6 +21,7 @@
     cmdline_run,
     settings,
 )
+from pydantic import Field
 
 YARA_EXTENSIONS = [".yar", ".yara"]
 
@@ -41,6 +42,8 @@ class AzulPluginYara(BinaryPlugin):
         size_before_disk=(int, 2**24),
         # Max number of yara includes to follow before giving up on looking for rules.
         max_yara_include_depth=(int, 5),
+        # Max number of yara streams to keep before dropping the rest.
+        max_yara_hit_streams_to_keep=(Annotated[int, Field(gt=0, le=100)], 50),
     )
 
     FEATURES = [
@@ -161,6 +164,8 @@ def execute(self, job: Job):
         # Read file from disk as multiple seek/read operations will be required
         fpath = job.get_data().get_filepath()
         found_raw_rule = dict()
+
+        yara_rule_streams_added = 0
         for match in matches.matching_rules:
             rule = match.namespace + "." + match.identifier
             self.add_feature_values("yararule", rule)
@@ -176,10 +181,13 @@ def execute(self, job: Job):
                 found_raw_rule[rule] = True
                 if new_rule not in seen_rules_md5s:
                     seen_rules_md5s.append(new_rule)
-                    raw_rule_with_header = (
-                        f"// plugin: {self.NAME}, namespace_identifier: {rule}\n".encode() + raw_rule
-                    )
-                    self.add_data(label=DataLabel.YARA_RULE_HIT, tags={}, data=raw_rule_with_header)
+                    # Add the original yara rule that hit as an augmented stream. Stop at max allowed Augmented streams.
+                    if yara_rule_streams_added < self.cfg.max_yara_hit_streams_to_keep:
+                        raw_rule_with_header = (
+                            f"// plugin: {self.NAME}, namespace_identifier: {rule}\n".encode() + raw_rule
+                        )
+                        self.add_data(label=DataLabel.YARA_RULE_HIT, tags={}, data=raw_rule_with_header)
+                        yara_rule_streams_added += 1
 
             for match_data in match.patterns:
                 var = match_data.identifier

tests/test_yara.py

@@ -1,5 +1,6 @@
+import tempfile
 import os
-
+from pydantic_core import ValidationError
 from azul_runner import FV, APIFeatureValue, Event, JobResult, State, test_template
 from azul_runner.models import EventData
 
@@ -296,6 +297,95 @@ def test_yara_match_result_disk(self):
             inspect_data=True,
         )
 
+    def test_yara_100_matching_rules(self):
+        """Ensure when there are a 100+ matching rules that ony the first 50 streams are kept.
+
+        This is because the binary ingestor can only process up to 100 streams before it rejects files anyway.
+        """
+        with tempfile.TemporaryDirectory("-yara-test") as temp_rule_dir:
+            for iteration in range(101):
+                raw_rule = (
+                    """rule Exploit_CVE_2015_0313_NewStream%d {
+    meta:
+        rule_group = "Exploit"  
+
+        //required
+        classification = "UNCLASSIFIED"
+        description = "Looks for presence of code that could indicate ANGLER EK use of this flash vuln"
+        exploit = "CVE-2015-0313"
+        info = "SWF"
+        organisation = "Defence"
+        poc = "azul@asd.gov.au" 
+        rule_version = "1"
+        yara_version = "1.6"
+
+        //optional
+        weight = 51
+
+    strings:
+        $ = "take_over_32("
+        $ = "get_x86_shellcode("
+        $ = "exploit_primordial_start("
+        $ = "exploit_primarodial_finish("
+        $ = "this.shellcodes.GetX86Shellcode("
+        $ = "Shellcodes("
+        $ = "attacking_buffer"
+        $ = "take_over_buffer"
+        $ = "make_spray_by_buffers_no_holes"
+        $ = "fake_object_address"
+    condition:
+	any of them
+}"""
+                    % iteration
+                )
+                with open(os.path.join(temp_rule_dir, f"ruleVersion-{iteration}.yara"), "w+") as f:
+                    f.write(raw_rule)
+
+            result = self.do_execution(
+                # This content should hit on the CVE-2015-0313 Angler EK rule in rules
+                data_in=[("content", b'example -> "exploit_primarodial_finish(" <-')],
+                config={
+                    "yara_rules_path": temp_rule_dir,
+                    "version_suffix": "0",
+                    "name_suffix": "0",
+                    "security_override": "OFFICIAL",
+                },
+            )
+            self.assertEqual(len(result.data.keys()), 50, "Must be capping the returned yara hits to 50")
+
+            # If option is overridden the number of streams kept should match that.
+            CONFIGURATION_OVERRIDE_FOR_MAX_YARA_STREAMS = 60
+            result = self.do_execution(
+                # This content should hit on the CVE-2015-0313 Angler EK rule in rules
+                data_in=[("content", b'example -> "exploit_primarodial_finish(" <-')],
+                config={
+                    "yara_rules_path": temp_rule_dir,
+                    "version_suffix": "0",
+                    "name_suffix": "0",
+                    "security_override": "OFFICIAL",
+                    "max_yara_hit_streams_to_keep": CONFIGURATION_OVERRIDE_FOR_MAX_YARA_STREAMS,
+                },
+            )
+            self.assertEqual(
+                len(result.data.keys()),
+                CONFIGURATION_OVERRIDE_FOR_MAX_YARA_STREAMS,
+                f"Must be capping the returned yara hits to {CONFIGURATION_OVERRIDE_FOR_MAX_YARA_STREAMS}",
+            )
+
+            # Fails when attempting to configure a value that must be less than 100.
+            with self.assertRaises(ValidationError):
+                result = self.do_execution(
+                    # This content should hit on the CVE-2015-0313 Angler EK rule in rules
+                    data_in=[("content", b'example -> "exploit_primarodial_finish(" <-')],
+                    config={
+                        "yara_rules_path": temp_rule_dir,
+                        "version_suffix": "0",
+                        "name_suffix": "0",
+                        "security_override": "OFFICIAL",
+                        "max_yara_hit_streams_to_keep": 200,
+                    },
+                )
+
     def test_yara_blacklist(self):
         """Blacklist should filter the only rule."""
         path = os.path.join(os.path.dirname(__file__), rel_rules_dir)