wip

Author: tomchop

core/database_arango.py

@@ -590,7 +590,7 @@ def link_to_tag(
             edge = json.loads(tag_relationship.model_dump_json())
             edge["_id"] = tag_relationship.id
             graph.update_edge(edge)
-            if self._collection_name != "auditlog":
+            if self._collection_name not in ("auditlog", "timeline"):
                 try:
                     event = message.TagEvent(
                         type=message.EventType.update,
@@ -627,7 +627,7 @@ def link_to_tag(
             time.sleep(ASYNC_JOB_WAIT_TIME)
         result = job.result()["new"]
         result["__id"] = result.pop("_key")
-        if self._collection_name != "auditlog":
+        if self._collection_name not in ("auditlog", "timeline"):
             try:
                 event = message.TagEvent(
                     type=message.EventType.new, tagged_object=self, tag_object=tag_obj
@@ -672,7 +672,7 @@ def clear_tags(self):
             time.sleep(ASYNC_JOB_WAIT_TIME)
         results = job.result()
         for edge in results["edges"]:
-            if self._collection_name != "auditlog":
+            if self._collection_name not in ("auditlog", "timeline"):
                 try:
                     job = self._db.collection("tagged").get(edge["_id"])
                     while job.status() != "done":
@@ -739,7 +739,7 @@ def link_to(
             job = async_graph.update_edge(edge)
             while job.status() != "done":
                 time.sleep(ASYNC_JOB_WAIT_TIME)
-            if self._collection_name != "auditlog":
+            if self._collection_name not in ("auditlog", "timeline"):
                 try:
                     event = message.LinkEvent(
                         type=message.EventType.update,
@@ -773,7 +773,7 @@ def link_to(
         result = job.result()["new"]
         result["__id"] = result.pop("_key")
         relationship = Relationship.load(result)
-        if self._collection_name != "auditlog":
+        if self._collection_name not in ("auditlog", "timeline"):
             try:
                 event = message.LinkEvent(
                     type=message.EventType.new,
@@ -910,7 +910,7 @@ def neighbors(
             str,
             "observable.ObservableTypes | entity.EntityTypes | indicator.IndicatorTypes | tag.Tag",
         ],
-        List[List["Relationship | TagRelationship"]],
+        List[List["RelationshipTypes"]],
         int,
     ]:
         """Fetches neighbors of the YetiObject.

core/events/consumers.py

@@ -4,6 +4,7 @@
 import logging
 import multiprocessing
 import os
+from collections import defaultdict
 
 from kombu import Connection, Exchange, Queue
 from kombu.mixins import ConsumerMixin
@@ -63,6 +64,9 @@ def get_consumers(self, consumer, channel):
             consumer(queues=self.queues, callbacks=[self.on_message], accept=["json"])
         ]
 
+MESSAGE_TYPES_COUNT = defaultdict(int)
+ctr = 0
+
 
 class EventConsumer(Consumer):
     def __init__(self, stop_event, connection, queues):
@@ -72,21 +76,29 @@ def debug(self, message, body):
         message_digest = hashlib.sha256(body.encode()).hexdigest()
         ts = int(message.timestamp.timestamp())
         if isinstance(message.event, ObjectEvent):
-            self.logger.debug(
+            self.logger.info(
                 f"Message received at {ts} - digest: {message_digest} | {message.event.event_message}"
             )
+            MESSAGE_TYPES_COUNT[message.event.event_message] += 1
         if isinstance(message.event, LinkEvent):
             source = message.event.link_source_event
             target = message.event.link_target_event
-            self.logger.debug(
+            self.logger.info(
                 f"Message received at {ts} - digest: {message_digest} | {source} --> {target}"
             )
+            MESSAGE_TYPES_COUNT[f"{source}->{target}"] += 1
         if isinstance(message.event, TagEvent):
-            self.logger.debug(
+            self.logger.info(
                 f"Message received at {ts} - digest: {message_digest} | {message.event.tag_message}"
             )
+            MESSAGE_TYPES_COUNT[message.event.tag_message] += 1
+
+        if ctr % 300 == 0:
+            for k, v in MESSAGE_TYPES_COUNT.items():
+                self.logger.info(f"{k}: {v}")
 
     def on_message(self, body, received_message):
+
         try:
             message = EventMessage(**json.loads(body))
         except Exception: