Ensure .data is always a dictionary/object when serializing events to JSON

bbot/core/event/base.py

@@ -374,6 +374,22 @@ def data_id(self):
     def _data_id(self):
         return self.data
 
+    def _json_data_key(self):
+        """
+        Because self.data can be either a string or a dictionary, this function is used to standardize
+        JSON-serialized events so that their .data attributes are _always_ dictionaries.
+
+        Events are serialized so that at the top level of their .data object there is a single
+        key which is their event type, and their data is the value.
+
+        For example, an `IP_ADDRESS` with .data of `"192.168.0.1"' will be `{"IP_ADDRESS": "192.168.0.1"}`.
+        A `DNS_NAME` of `evilcorp.com` will be `{"DNS_NAME": "evilcorp.com"}`, etc.
+
+        Returns:
+            The key to be used when constructing a dictionary of the event's data.
+        """
+        return self.type
+
     @property
     def pretty_string(self):
         """
@@ -431,9 +447,9 @@ def json(self, mode="json"):
                 j.update({i: v})
         data_attr = getattr(self, f"data_{mode}", None)
         if data_attr is not None:
-            j["data"] = data_attr
+            j["data"] = {self._json_data_key(): data_attr}
         else:
-            j["data"] = smart_decode(self.data)
+            j["data"] = {self._json_data_key(): smart_decode(self.data)}
         web_spider_distance = getattr(self, "web_spider_distance", None)
         if web_spider_distance is not None:
             j["web_spider_distance"] = web_spider_distance
@@ -1037,10 +1053,11 @@ def make_event(
 
 
 def event_from_json(j):
+    event_type = j["type"]
     try:
         kwargs = {
-            "data": j["data"],
-            "event_type": j["type"],
+            "data": j["data"][event_type],
+            "event_type": event_type,
             "scans": j.get("scans", []),
             "tags": j.get("tags", []),
             "confidence": j.get("confidence", 5),

bbot/test/test_step_1/test_events.py

@@ -342,7 +342,7 @@ async def test_events(events, scan, helpers, bbot_config):
     timestamp = db_event.timestamp.timestamp()
     json_event = db_event.json()
     assert json_event["scope_distance"] == 1
-    assert json_event["data"] == "127.0.0.1"
+    assert json_event["data"] == {"IP_ADDRESS": "127.0.0.1"}
     assert json_event["type"] == "IP_ADDRESS"
     assert json_event["timestamp"] == timestamp
     reconstituted_event = event_from_json(json_event)
@@ -355,7 +355,8 @@ async def test_events(events, scan, helpers, bbot_config):
     assert http_response.source_id == scan.root_event.id
     assert http_response.data["input"] == "http://example.com:80"
     json_event = http_response.json(mode="graph")
-    assert isinstance(json_event["data"], str)
+    assert isinstance(json_event["data"], dict)
+    assert isinstance(json_event["data"]["HTTP_RESPONSE"], str)
     json_event = http_response.json()
     assert isinstance(json_event["data"], dict)
     assert json_event["type"] == "HTTP_RESPONSE"

bbot/test/test_step_2/module_tests/test_module_json.py

@@ -9,6 +9,9 @@ def check(self, module_test, events):
         txt_file = module_test.scan.home / "output.ndjson"
         lines = list(module_test.scan.helpers.read_file(txt_file))
         assert lines
-        e = event_from_json(json.loads(lines[0]))
+        json_event = json.loads(lines[0])
+        assert json_event["type"] == "SCAN"
+        assert json_event["data"] == {"SCAN": f"{module_test.scan.name} ({module_test.scan.id})"}
+        e = event_from_json(json_event)
         assert e.type == "SCAN"
         assert e.data == f"{module_test.scan.name} ({module_test.scan.id})"