Regex matching in /graph/match (#1112)

Co-authored-by: Sebastien Larinier <sebdraven@protonmail.com>

core/database_arango.py

@@ -759,6 +759,19 @@ def filter(
                 conditions.append(f"o.@arg{i}_key IN @arg{i}_value")
                 aql_args[f"arg{i}_key"] = key[:-4]
                 sorts.append(f"o.@arg{i}_key")
+            elif key.endswith("__in~"):
+                del aql_args[f"arg{i}_value"]
+                if not value:
+                    continue
+                aql_args[f"arg{i}_key"] = key[:-5]
+                or_conditions = []
+                for j, v in enumerate(value):
+                    or_conditions.append(
+                        f"REGEX_TEST(o.@arg{i}_key, @arg{i}{j}_value, true)"
+                    )
+                    aql_args[f"arg{i}{j}_value"] = v.strip()
+                    sorts.append(f"o.@arg{i}_key")
+                conditions.append(f"({' OR '.join(or_conditions)})")
             elif key in ["labels", "relevant_tags"]:
                 conditions.append(f"@arg{i}_value ALL IN o.@arg{i}_key")
                 aql_args[f"arg{i}_key"] = key

core/web/apiv2/graph.py

@@ -214,6 +214,7 @@ async def delete(relationship_id: str) -> None:
 class AnalysisRequest(BaseModel):
     observables: list[str]
     add_tags: list[str] = []
+    regex_match: bool = False
     add_type: observable.ObservableType | None = None
     fetch_neighbors: bool = True
     add_unknown: bool = False
@@ -232,6 +233,7 @@ async def match(request: AnalysisRequest) -> AnalysisResponse:
     """Fetches neighbors for a given Yeti Object."""
 
     entities = []  # type: list[tuple[graph.Relationship, entity.Entity]]
+    seen_entities = set()
     observables = []  # type: list[tuple[graph.Relationship, observable.Observable]]
 
     unknown = set(request.observables)
@@ -249,8 +251,11 @@ async def match(request: AnalysisRequest) -> AnalysisResponse:
 
             unknown.discard(value)
 
+    operator = "value__in"
+    if request.regex_match:
+        operator = "value__in~"
     db_observables, _ = observable.Observable.filter(
-        query_args={"value__in": request.observables},
+        query_args={operator: request.observables},
         graph_queries=[("tags", "tagged", "outbound", "name")],
     )
     for db_observable in db_observables:
@@ -271,7 +276,10 @@ async def match(request: AnalysisRequest) -> AnalysisResponse:
                     other = vertices[edge.target]
 
                 if isinstance(other, entity.Entity):
+                    if other.extended_id in seen_entities:
+                        continue
                     entities.append((edge, other))
+                    seen_entities.add(other.extended_id)
                 if isinstance(other, observable.Observable):
                     observables.append((edge, other))
 

tests/apiv2/graph.py

@@ -366,6 +366,28 @@ def test_existing_links(self):
         self.assertEqual(entity["type"], "threat-actor")
         self.assertEqual(entity["name"], "tester")
 
+    def test_match(self):
+        response = client.post(
+            "/api/v2/graph/match",
+            json={"observables": ["test1.com"]},
+        )
+        data = response.json()
+        self.assertEqual(response.status_code, 200, data)
+        self.assertEqual(len(data["known"]), 1)
+        self.assertEqual(data["known"][0]["value"], "test1.com")
+
+    def test_match_regex(self):
+        response = client.post(
+            "/api/v2/graph/match",
+            json={"observables": ["tes.[0-9]+.com"], "regex_match": True},
+        )
+        data = response.json()
+        self.assertEqual(response.status_code, 200, data)
+        self.assertEqual(len(data["known"]), 3)
+        self.assertEqual(data["known"][0]["value"], "http://test1.com/admin")
+        self.assertEqual(data["known"][1]["value"], "test1.com")
+        self.assertEqual(data["known"][2]["value"], "test2.com")
+
     def test_matches_exist(self):
         """Tests that indicator matches will surface."""
         response = client.post(