Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Regex matching in /graph/match #1112

Merged
merged 11 commits into from
Jul 23, 2024
13 changes: 13 additions & 0 deletions core/database_arango.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
10 changes: 9 additions & 1 deletion core/web/apiv2/graph.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)
Expand All @@ -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:
Expand All @@ -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))

Expand Down
22 changes: 22 additions & 0 deletions tests/apiv2/graph.py
tomchop marked this conversation as resolved.
Show resolved Hide resolved
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
Loading