feat: add option to exclude raw rules upload (#52)

feat: add option to exclude raw rules upload

Author: kevin-acsc

azul_plugin_yara/main.py

@@ -44,6 +44,7 @@ class AzulPluginYara(BinaryPlugin):
         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),
+        yara_upload_raw_rules_as_streams=(bool, True),
     )
 
     FEATURES = [
@@ -175,23 +176,24 @@ def execute(self, job: Job):
             self.add_feature_values("yararule", rule)
             found_raw_rule[rule] = False
 
-            # Find the raw rule and save it as a file
-            rule_file_path = self.namespace_to_rule_path[match.namespace]
-            self.yara_include_depth = 0
-            raw_rule = self.fetch_original_rule(rule_file_path, match.identifier, self.logger)
-            if len(raw_rule) > 0:
-                new_rule = md5(raw_rule).hexdigest()  # noqa: S324
-                # Rule was found
-                found_raw_rule[rule] = True
-                if new_rule not in seen_rules_md5s:
-                    seen_rules_md5s.append(new_rule)
-                    # 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:  # ty: ignore[unresolved-attribute] ty doesn't understand add_settings
-                        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
+            if self.cfg.yara_upload_raw_rules_as_streams:  # ty: ignore[unresolved-attribute] ty doesn't understand add_settings
+                # Find the raw rule and save it as a file
+                rule_file_path = self.namespace_to_rule_path[match.namespace]
+                self.yara_include_depth = 0
+                raw_rule = self.fetch_original_rule(rule_file_path, match.identifier, self.logger)
+                if len(raw_rule) > 0:
+                    new_rule = md5(raw_rule).hexdigest()  # noqa: S324
+                    # Rule was found
+                    found_raw_rule[rule] = True
+                    if new_rule not in seen_rules_md5s:
+                        seen_rules_md5s.append(new_rule)
+                        # 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:  # ty: ignore[unresolved-attribute] ty doesn't understand add_settings
+                            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

@@ -386,6 +386,80 @@ def test_yara_100_matching_rules(self):
                     },
                 )
 
+    def test_yara_no_rule_stream_upload(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.
+            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",
+                    "yara_upload_raw_rules_as_streams": "False",
+                },
+            )
+            self.assertEqual(
+                len(result.data.keys()),
+                0,
+                f"Must not upload any rules as streams if yara_upload_raw_rules_as_streams==False",
+            )
+
     def test_yara_blacklist(self):
         """Blacklist should filter the only rule."""
         path = os.path.join(os.path.dirname(__file__), rel_rules_dir)