apply event .data change across all event types

blacklanternsecurity · Sep 22, 2023 · 637bd90 · 637bd90
1 parent 1b53432
commit 637bd90
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 10 deletions.
diff --git a/bbot/core/event/base.py b/bbot/core/event/base.py
@@ -446,12 +446,10 @@ def json(self, mode="json"):
             if v:
                 j.update({i: v})
         data_attr = getattr(self, f"data_{mode}", None)
-        if isinstance(data_attr, str):
-            data_attr = {self._json_data_key(): data_attr}
         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
@@ -1055,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),

diff --git a/bbot/test/test_step_1/test_events.py b/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"

diff --git a/bbot/test/test_step_2/module_tests/test_module_json.py b/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 == {"SCAN": f"{module_test.scan.name} ({module_test.scan.id})"}
+        assert e.data == f"{module_test.scan.name} ({module_test.scan.id})"