Yara bundles & optimized validation (#1230)

Author: tomchop

core/errors.py

@@ -1,6 +1,6 @@
 class YetiError(RuntimeError):
-    def __init__(self, message: str, meta: dict):
-        self.meta = meta
+    def __init__(self, message: str, meta: dict | None = None):
+        self.meta = meta or {}
         super().__init__(message)
 
 

core/schemas/indicators/yara.py

@@ -1,4 +1,5 @@
-from typing import Any, ClassVar, Literal
+import logging
+from typing import ClassVar, Literal
 
 import plyara
 import plyara.exceptions
@@ -19,6 +20,9 @@
 }
 
 
+logger = logging.getLogger(__name__)
+
+
 class MatchInstance(BaseModel):
     """
     Represents an instance of a string match.
@@ -122,10 +126,8 @@ class Yara(indicator.Indicator):
     dependencies: list[str] = []
     private: bool = False
 
-    @model_validator(mode="before")
-    @classmethod
-    def validate_yara(cls, data: Any):
-        rule = data.get("pattern")
+    def validate_yara(self):
+        rule = self.pattern
         if not rule:
             raise ValueError("Yara rule body is required.")
         try:
@@ -138,13 +140,16 @@ def validate_yara(cls, data: Any):
             raise ValueError("No valid Yara rules found in the rule body.")
         parsed_rule = rules[0]
         rule_deps = set(plyara.utils.detect_dependencies(parsed_rule))
-        data["dependencies"] = rule_deps - ALLOWED_EXTERNALS.keys()
-        data["name"] = parsed_rule["rule_name"]
-        data["private"] = "private" in parsed_rule.get("scopes", [])
-
-        return data
+        self.dependencies = list(rule_deps - ALLOWED_EXTERNALS.keys())
+        self.name = parsed_rule["rule_name"]
+        self.private = "private" in parsed_rule.get("scopes", [])
 
     def save(self):
+        try:
+            self.validate_yara()
+        except ValueError as error:
+            raise errors.ObjectCreationError(str(error)) from error
+
         missing_deps = []
         for dep_name in self.dependencies:
             dep = Yara.find(name=dep_name)
@@ -247,19 +252,36 @@ def rule_with_dependencies(
 
         concatenated_rules = ""
 
-        parsed_rule = plyara.Plyara().parse_string(self.pattern)[0]
-        dependencies = plyara.utils.detect_dependencies(parsed_rule)
+        dependencies = self.dependencies
 
         for dependency in dependencies:
+            if dependency in resolved:
+                continue
+            logger.info(f"Resolving dependency: {dependency}")
             dep_rule = Yara.find(name=dependency)
             if not dep_rule:
                 raise ValueError(f"Rule depends on unknown dependency '{dependency}'")
-            if dep_rule.name not in resolved:
-                concatenated_rules += dep_rule.rule_with_dependencies(resolved, seen)
+            concatenated_rules += dep_rule.rule_with_dependencies(resolved, seen)
 
         if self.name not in resolved:
             concatenated_rules += self.pattern + "\n\n"
             resolved.add(self.name)
 
         seen.remove(self.name)
         return concatenated_rules
+
+    @classmethod
+    def generate_yara_bundle(cls, rules: list["Yara"]) -> str:
+        """Export a list of Yara rules to a single string.
+
+        Args:
+            rules: A list of Yara rules to export.
+
+        Returns:
+            A string containing the exported rules.
+        """
+        resolved: set[str] = set()
+        bulk_rules = ""
+        for rule in rules:
+            bulk_rules += rule.rule_with_dependencies(resolved)
+        return bulk_rules

core/web/apiv2/indicators.py

@@ -1,3 +1,5 @@
+import logging
+
 from fastapi import APIRouter, HTTPException, Request
 from pydantic import BaseModel, ConfigDict, Field, conlist
 
@@ -8,10 +10,13 @@
     Indicator,
     IndicatorType,
     IndicatorTypes,
+    Yara,
 )
 from core.schemas.rbac import global_permission, permission_on_ids, permission_on_target
 from core.schemas.tag import MAX_TAGS_REQUEST
 
+logger = logging.getLogger(__name__)
+
 
 # Request schemas
 class NewIndicatorRequest(BaseModel):
@@ -59,6 +64,20 @@ class IndicatorTagResponse(BaseModel):
     tags: dict[str, dict[str, graph.TagRelationship]]
 
 
+class YaraBundleRequest(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    ids: list[str] = []
+    tags: list[str] = []
+    exclude_tags: list[str] = []
+
+
+class YaraBundleResponse(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    bundle: str
+
+
 # API endpoints
 router = APIRouter()
 
@@ -171,3 +190,34 @@ def tag(httpreq: Request, request: IndicatorTagRequest) -> IndicatorTagResponse:
         indicator_tags[db_indicator.extended_id] = db_indicator.tags
 
     return IndicatorTagResponse(tagged=len(indicators), tags=indicator_tags)
+
+
+@router.post("/yara/bundle")
+def get_yara_bundle(httpreq: Request, request: YaraBundleRequest) -> YaraBundleResponse:
+    """Generates a YARA bundle from a list of indicators."""
+    indicators = []
+    for indicator_id in request.ids:
+        db_indicator = Yara.get(indicator_id)
+        if not db_indicator:
+            raise HTTPException(
+                status_code=404,
+                detail=f"YARA bundle request contained an unknown indicator: ID:{indicator_id}",
+            )
+        indicators.append(db_indicator)
+    import time
+
+    indicators_from_tags, _ = Indicator.filter(
+        query_args={"type": "yara"},
+        tag_filter=request.tags,
+        graph_queries=[("tags", "tagged", "outbound", "name")],
+        user=httpreq.state.user,
+    )
+
+    for indicator in indicators_from_tags:
+        if any(tag in request.exclude_tags for tag in indicator.tags):
+            continue
+        indicators.append(indicator)
+
+    bundle = Yara.generate_yara_bundle(rules=indicators)
+
+    return YaraBundleResponse(bundle=bundle)

tests/apiv2/indicators.py

@@ -172,10 +172,10 @@ def test_bad_yara(self):
             "/api/v2/indicators/",
             json={"indicator": indicator_dict},
         )
-        self.assertEqual(response.status_code, 422)
+        self.assertEqual(response.status_code, 400)
         data = response.json()
         self.assertIn(
-            "No valid Yara rules found in the rule body", data["detail"][0]["msg"]
+            "No valid Yara rules found in the rule body", data["detail"]["description"]
         )
 
     def test_bad_yara_graceful_failure(self):

tests/schemas/yararule.py

@@ -1,6 +1,6 @@
 import unittest
 
-from core import database_arango
+from core import database_arango, errors
 from core.schemas.indicator import DiamondModel
 from core.schemas.indicators.yara import Yara
 
@@ -29,12 +29,13 @@ def test_yara_name_and_deps(self):
             location="any",
             diamond=DiamondModel.capability,
         )
+        yara.validate_yara()
 
         self.assertEqual(yara.name, "test")
         self.assertEqual(yara.dependencies, ["dep"])
 
     def test_invalid_yara_rule(self):
-        with self.assertRaises(ValueError) as error:
+        with self.assertRaises(errors.ObjectCreationError) as error:
             Yara(
                 pattern='rule test { wooo: $a = "test" fooo: $a and dep }',
                 location="any",
@@ -44,7 +45,7 @@ def test_invalid_yara_rule(self):
         self.assertIn("Unknown text wooo", str(error.exception))
 
     def test_fail_on_more_than_one_rule(self):
-        with self.assertRaises(ValueError) as error:
+        with self.assertRaises(errors.ObjectCreationError) as error:
             Yara(
                 pattern="rule test { condition: true } rule test2 { condition: true }",
                 location="any",
@@ -91,6 +92,47 @@ def test_dependency_calculation(self):
             ),
         )
 
+    def test_bulk_dependency_export(self):
+        Yara(
+            pattern="rule dep0 { condition: true }",
+            location="any",
+            diamond=DiamondModel.capability,
+        ).save()
+
+        Yara(
+            pattern="rule dep1 { condition: true and dep0 }",
+            location="any",
+            diamond=DiamondModel.capability,
+        ).save()
+
+        Yara(
+            pattern="rule dep2 { condition: true and dep1 }",
+            location="any",
+            diamond=DiamondModel.capability,
+        ).save()
+
+        yara_rule = Yara(
+            pattern="rule test { condition: true and dep2 and dep1 }",
+            location="any",
+            diamond=DiamondModel.capability,
+        ).save()
+
+        yara_rule2 = Yara(
+            pattern="rule test2 { condition: true and dep2 }",
+            location="any",
+            diamond=DiamondModel.capability,
+        )
+
+        export = Yara.generate_yara_bundle([yara_rule, yara_rule2])
+        self.assertEqual(
+            export,
+            "rule dep0 { condition: true }\n\n"
+            "rule dep1 { condition: true and dep0 }\n\n"
+            "rule dep2 { condition: true and dep1 }\n\n"
+            "rule test { condition: true and dep2 and dep1 }\n\n"
+            "rule test2 { condition: true and dep2 }\n\n",
+        )
+
     def test_yara_dependency_creates_links(self):
         dep0 = Yara(
             pattern="rule dep0 { condition: true }",