Feat: Add ability to search by date (#908)

Author: tomchop

core/database_arango.py

@@ -670,6 +670,14 @@ def filter(
                 conditions.append(f"COUNT(FOR c IN o.context[*] FILTER REGEX_TEST(c.@arg{i}_key, @arg{i}_value) RETURN c) > 0")
                 aql_args[f'arg{i}_key'] = context_field
                 sorts.append(f"o.context[*].@arg{i}_key")
+            elif key == 'created':
+                operator = value[0]
+                if operator not in ['<', '>']:
+                    operator = '='
+                else:
+                    aql_args[f'arg{i}_value'] = value[1:]
+                conditions.append(f"DATE_TIMESTAMP(o.created) {operator}= DATE_TIMESTAMP(@arg{i}_value)")
+                sorts.append("o.created")
             else:
                 conditions.append(f"REGEX_TEST(o.@arg{i}_key, @arg{i}_value, true)")
                 aql_args[f'arg{i}_key'] = key

core/schemas/entity.py

@@ -1,14 +1,13 @@
 import datetime
 import re
-import unicodedata
 from enum import Enum
-from typing import ClassVar, Literal, Optional, Type
+from typing import ClassVar, Literal, Type
+
+from pydantic import BaseModel, Field
 
 from core import database_arango
 from core.helpers import now
 from core.schemas.graph import TagRelationship
-from core.schemas.tag import Tag
-from pydantic import BaseModel, Field
 
 
 class EntityType(str, Enum):

tests/apiv2/entities.py

@@ -1,3 +1,4 @@
+import datetime
 import unittest
 
 from fastapi.testclient import TestClient
@@ -12,7 +13,10 @@
 class EntityTest(unittest.TestCase):
     def setUp(self) -> None:
         database_arango.db.clear()
-        self.entity1 = entity.ThreatActor(name="ta1", aliases=["badactor"]).save()
+        self.entity1 = entity.ThreatActor(
+            name="ta1",
+            aliases=["badactor"],
+            created=datetime.datetime(2020, 1, 1)).save()
         self.entity1.tag(["ta1"])
         self.entity2 = entity.ThreatActor(name="bears").save()
 
@@ -66,6 +70,35 @@ def test_search_entities_subfields(self):
         self.assertEqual(data["entities"][0]["name"], "ta1")
         self.assertEqual(data["entities"][0]["type"], "threat-actor")
 
+    def test_search_entities_with_creation_date(self):
+        response = client.post(
+            "/api/v2/entities/search",
+            json={
+                "query": {"created": "<2020-01-02"},
+                "type": "threat-actor",
+            },
+        )
+        data = response.json()
+        self.assertEqual(response.status_code, 200, data)
+        self.assertEqual(len(data["entities"]), 1)
+        self.assertEqual(data["entities"][0]["name"], "ta1")
+        self.assertEqual(data["entities"][0]["type"], "threat-actor")
+
+        response = client.post(
+            "/api/v2/entities/search",
+            json={
+                "query": {"created": ">2020-01-03"},
+                "type": "threat-actor",
+            },
+        )
+        data = response.json()
+        self.assertEqual(response.status_code, 200, data)
+        # self.entity2 is created with a default timestamp of now(), so should
+        # be the only threat-actor entity captured with this filter.
+        self.assertEqual(len(data["entities"]), 1, data)
+        self.assertEqual(data["entities"][0]["name"], "bears")
+        self.assertEqual(data["entities"][0]["type"], "threat-actor")
+
     def test_new_entity_with_tag(self):
         response = client.post(
             "/api/v2/entities/",

tests/schemas/entity.py

@@ -5,15 +5,15 @@
 from core.schemas.entity import (AttackPattern, Entity, Malware, ThreatActor,
                                  Tool, Vulnerability)
 from core.schemas.observables import hostname
-
+import datetime
 
 class EntityTest(unittest.TestCase):
 
     def setUp(self) -> None:
         database_arango.db.clear()
         self.ta1 = ThreatActor(name="APT123", aliases=['CrazyFrog']).save()
         self.vuln1 = Vulnerability(name="CVE-2018-1337", title='elite exploit').save()
-        self.malware1 = Malware(name="zeus").save()
+        self.malware1 = Malware(name="zeus", created=datetime.datetime(2020, 1, 1)).save()
         self.tool1 = Tool(name="mimikatz").save()
 
     def tearDown(self) -> None:
@@ -70,6 +70,22 @@ def test_filter_entities_regex(self):
         self.assertEqual(total, 1)
         self.assertEqual(entities[0], self.vuln1)
 
+    def test_filter_entities_time(self):
+        entities, total = Entity.filter({"created": "2020-01-01"})
+        self.assertEqual(len(entities), 1)
+        self.assertEqual(total, 1)
+        self.assertEqual(entities[0], self.malware1)
+
+        entities, total = Entity.filter({"created": "<2020-01-02"})
+        self.assertEqual(len(entities), 1)
+        self.assertEqual(total, 1)
+        self.assertEqual(entities[0], self.malware1)
+
+        entities, total = Entity.filter({"created": ">2020-01-02"})
+        self.assertEqual(len(entities), 3)
+        self.assertEqual(total, 3)
+        self.assertNotIn(self.malware1, entities)
+
     def test_entity_with_tags(self):
         entity = ThreatActor(name="APT0").save()
         entity.tag(['tag1', 'tag2'])