Global search (#1257)

Author: tomchop

core/database_arango.py

@@ -180,8 +180,7 @@ def create_analyzers(self):
         self.db.create_analyzer(
             name="norm",
             analyzer_type="norm",
-            properties={"locale": "en", "accent": False, "case": "lower"},
-            features=[],
+            properties={"locale": "en.utf-8", "accent": False, "case": "lower"},
         )
 
     def refresh_views(self):
@@ -306,7 +305,6 @@ def create_views(self):
             link_definitions[view_target] = {
                 "analyzers": ["identity", "norm"],
                 "includeAllFields": True,
-                "storedValues": [{"fields": ["name", "tags", "type"]}],
                 "trackListPositions": False,
             }
 
@@ -333,6 +331,20 @@ def create_views(self):
         except Exception:
             pass
 
+        for target in link_definitions:
+            del link_definitions[target]["analyzers"]
+            link_definitions[target]["analyzers"] = []
+            link_definitions[target]["includeAllFields"] = False
+            link_definitions[target]["fields"] = {
+                "tags": {"fields": {"name": {"analyzers": ["identity", "norm"]}}},
+                "dfiq_tags": {"analyzers": ["identity", "norm"]},
+                "type": {"analyzers": ["identity", "norm"]},
+                "root_type": {"analyzers": ["identity", "norm"]},
+                "value": {"analyzers": ["identity", "norm"]},
+                "name": {"analyzers": ["identity", "norm"]},
+                "created": {"analyzers": ["identity", "norm"]},
+            }
+
         self.db.create_arangosearch_view(
             name="all_objects_view",
             properties={
@@ -343,6 +355,7 @@ def create_views(self):
                     {"field": "created", "direction": "desc"},
                     {"field": "value", "direction": "asc"},
                     {"field": "name", "direction": "asc"},
+                    {"field": "tags.name", "direction": "asc"},
                 ],
             },
         )
@@ -439,6 +452,7 @@ class ArangoYetiConnector(AbstractYetiConnector):
     """Yeti connector for an ArangoDB backend."""
 
     _db = db
+    _collection_name: str | None = None
 
     def __init__(self):
         self._arango_id = None
@@ -1016,6 +1030,7 @@ def filter(
             offset: Skip this many objects when querying the DB.
             count: How many objecst after `offset` to return.
             sorting: A list of (order, ascending) fields to sort by.
+            aliases: A list of (alias, type) tuples to use for filtering.
             graph_queries: A list of (name, graph, direction, field) tuples to
                 query the graph with.
             wildcard: whether all values should be interpreted as wildcard searches.
@@ -1026,11 +1041,19 @@ def filter(
         """
         cls._get_collection()
         colname = cls._collection_name
+        if colname is None:
+            colname = "all_objects_view"
         conditions = []
         filter_conditions = []  # used for clauses that are not supported by arangosearch
         sorts = []
 
         using_view = False
+        generic_query = False
+
+        if colname == "all_objects_view":
+            generic_query = True
+            using_view = True
+
         if (
             query_args
             and colname in ("observables", "entities", "indicators", "dfiq")
@@ -1085,9 +1108,13 @@ def filter(
                 aql_args[f"arg{i}_key"] = key
             elif key == "tags":
                 if using_view:
-                    conditions.append(f"@arg{i}_value ALL IN o.tags.name")
+                    conditions.append(
+                        f"(FOR t in @arg{i}_value RETURN LOWER(t)) ALL IN o.tags.name"
+                    )
                 else:
-                    conditions.append(f"@arg{i}_value ALL IN o.tags[*].name")
+                    conditions.append(
+                        f"(FOR t in @arg{i}_value RETURN LOWER(t)) ALL IN o.tags[*].name"
+                    )
             elif key in ("created", "modified", "tags.expires"):
                 # Value is a string, we're checking the first character.
                 operator = value[0]
@@ -1114,11 +1141,16 @@ def filter(
                     key_conditions = [f"REGEX_TEST(o.@arg{i}_key, @arg{i}_value, true)"]
 
                 for alias, alias_type in aliases:
-                    if (
-                        alias_type in {"text", "option"}
-                        or alias_type == "list"
-                        and using_view
-                    ):
+                    if alias == "tags":
+                        if using_view:
+                            key_conditions.append(
+                                f"ANALYZER(LIKE(o.tags.name, LOWER(@arg{i}_value)), 'norm')"
+                            )
+                        else:
+                            key_conditions.append(
+                                f"LOWER(@arg{i}_value) IN o.tags[*].name"
+                            )
+                    if alias_type in {"text", "option", "list"} and using_view:
                         if using_view and not using_regex:
                             key_conditions.append(
                                 f"ANALYZER(LIKE(o.{alias}, LOWER(@arg{i}_value)), 'norm')"
@@ -1232,7 +1264,14 @@ def filter(
         results = []
         for doc in documents:
             doc["__id"] = doc.pop("_key")
-            results.append(cls.load(doc))
+            if not generic_query:
+                results.append(cls.load(doc))
+            else:
+                # Generic objects are not loaded, they are returned as dicts.
+                doc["id"] = doc.pop("__id")
+                del doc["_id"]
+                del doc["_rev"]
+                results.append(doc)
         total = stats.get("fullCount", len(results))
         return results, total or 0
 
@@ -1321,7 +1360,10 @@ def _get_collection(cls):
         Returns:
           The ArangoDB collection corresponding to the object class.
         """
-        return cls._db.collection(cls._collection_name)
+        if cls._collection_name is not None:
+            return cls._db.collection(cls._collection_name)
+        else:
+            return "all_objects_view"
 
 
 def tagged_observables_export(cls, args):

core/web/apiv2/search.py

@@ -0,0 +1,41 @@
+from fastapi import APIRouter, Request
+from pydantic import BaseModel, ConfigDict
+
+from core.database_arango import ArangoYetiConnector
+
+# API endpoints
+router = APIRouter()
+
+
+class SearchRequest(BaseModel):
+    """Search request message."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    query: dict[str, str | int | list] = {}
+    sorting: list[tuple[str, bool]] = []
+    filter_aliases: list[tuple[str, str]] = []
+    count: int = 50
+    page: int = 0
+
+
+class SearchResponse(BaseModel):
+    """Search response message."""
+
+    results: list[dict]
+    total: int = 0
+
+
+@router.post("/")
+def search(httpreq: Request, request: SearchRequest) -> SearchResponse:
+    """Gets the system config."""
+    results, total = ArangoYetiConnector.filter(
+        request.query,
+        sorting=request.sorting,
+        aliases=request.filter_aliases,
+        count=request.count,
+        offset=request.page * request.count,
+        links_count=True,
+        user=httpreq.state.user,
+    )
+    return SearchResponse(results=results, total=total)

core/web/webapp.py

@@ -19,6 +19,7 @@
     indicators,
     observables,
     rbac,
+    search,
     system,
     tag,
     tasks,
@@ -40,6 +41,13 @@
 
 api_router.include_router(audit.router, prefix="/audit", tags=["audit"])
 
+api_router.include_router(
+    search.router,
+    prefix="/search",
+    tags=["search"],
+    dependencies=[Depends(auth.get_current_active_user)],
+)
+
 api_router.include_router(
     observables.router,
     prefix="/observables",

tests/apiv2/search.py

@@ -0,0 +1,130 @@
+import logging
+import sys
+import time
+import unittest
+
+from fastapi.testclient import TestClient
+
+from core import database_arango
+from core.schemas import dfiq, entity, indicator, observable
+from core.schemas.user import UserSensitive
+from core.web import webapp
+
+client = TestClient(webapp.app)
+
+
+class searchTest(unittest.TestCase):
+    def setUp(self) -> None:
+        logging.disable(sys.maxsize)
+        database_arango.db.connect(database="yeti_test")
+        database_arango.db.truncate()
+
+        user = UserSensitive(username="test", password="test", enabled=True).save()
+        apikey = user.create_api_key("default")
+        token_data = client.post(
+            "/api/v2/auth/api-token", headers={"x-yeti-apikey": apikey}
+        ).json()
+        client.headers = {"Authorization": "Bearer " + token_data["access_token"]}
+
+        entity.Malware(
+            name="test_malware",
+            description="Test malware entity",
+            type="malware",
+        ).save()
+        m2 = entity.Malware(
+            name="tagged_malware",
+            description="malware entity 2",
+            type="malware",
+        ).save()
+        m2.tag("tagged")
+        m2.tag("global")
+
+        r = indicator.Regex(
+            name="test_regex_global",
+            description="Test regex indicator",
+            type="regex",
+            pattern="^test.*",
+            diamond="victim",
+        ).save()
+        o = observable.Hostname(
+            description="Test hostname observable",
+            type="hostname",
+            value="test.tomchop.me",
+        ).save()
+        dfiq.DFIQScenario(
+            name="test_dfiq",
+            description="Test DFIQ",
+            dfiq_tags=["tagged", "global"],
+            dfiq_version="1.0.1",
+            dfiq_yaml="name: test_dfiq\nversion: 1.0.1\ndescription: Test DFIQ",
+        ).save()
+        r.tag("tagged")
+        o.tag(["tagged", "global"])
+        time.sleep(5)
+
+    def test_search_name_or_value(self) -> None:
+        """Test global search by name or value."""
+        params = {"query": {"name": "test"}, "filter_aliases": [["value", "text"]]}
+        response = client.post("/api/v2/search", json=params)
+        self.assertEqual(response.status_code, 200, response.text)
+        data = response.json()
+        self.assertEqual(response.status_code, 200, data)
+        self.assertIn("results", data)
+        self.assertEqual(len(data["results"]), 4, data)
+        names_or_values = [
+            r["name"] if "name" in r else r["value"] for r in data["results"]
+        ]
+        self.assertCountEqual(
+            names_or_values,
+            ["test_malware", "test_regex_global", "test.tomchop.me", "test_dfiq"],
+            data,
+        )
+
+    def test_search_tag(self) -> None:
+        """Test global search by tag."""
+        params = {
+            "query": {"tags": ["tagged"]},
+        }
+        response = client.post("/api/v2/search", json=params)
+        self.assertEqual(response.status_code, 200, response.text)
+        data = response.json()
+        self.assertEqual(response.status_code, 200, data)
+        self.assertIn("results", data)
+        self.assertEqual(len(data["results"]), 3, data)
+        names_or_values = [
+            r["name"] if "name" in r else r["value"] for r in data["results"]
+        ]
+        self.assertCountEqual(
+            names_or_values,
+            [
+                "test_regex_global",
+                "test.tomchop.me",
+                "tagged_malware",
+            ],
+            data,
+        )
+
+    def test_search_more_fields(self) -> None:
+        """Test global search with more fields."""
+        params = {
+            "query": {"name": "global"},
+            "filter_aliases": [
+                ["value", "text"],
+                ["tags", ""],
+                ["dfiq_tags", "list"],
+            ],
+        }
+        response = client.post("/api/v2/search", json=params)
+        self.assertEqual(response.status_code, 200, response.text)
+        data = response.json()
+        self.assertEqual(response.status_code, 200, data)
+        self.assertIn("results", data)
+        self.assertEqual(len(data["results"]), 4, data)
+        names_or_values = [
+            r["name"] if "name" in r else r["value"] for r in data["results"]
+        ]
+        self.assertCountEqual(
+            names_or_values,
+            ["test_regex_global", "test.tomchop.me", "tagged_malware", "test_dfiq"],
+            data,
+        )