Pydantic updates

bbot/models/pydantic.py

@@ -1,7 +1,7 @@
 import json
 import logging
-from typing import Optional, List, Annotated
-from pydantic import BaseModel, ConfigDict, Field
+from pydantic import BaseModel, ConfigDict, Field, computed_field
+from typing import Optional, List, Annotated, get_origin, get_args
 
 from bbot.models.helpers import utc_now_timestamp
 
@@ -11,11 +11,6 @@
 class BBOTBaseModel(BaseModel):
     model_config = ConfigDict(extra="ignore")
 
-    def __init__(self, **data):
-        super().__init__(**data)
-        if getattr(self, "host", ""):
-            self.reverse_host = self.host[::-1]
-
     def to_json(self, **kwargs):
         return json.dumps(self.model_dump(), sort_keys=True, **kwargs)
 
@@ -26,8 +21,24 @@ def __eq__(self, other):
         return hash(self) == hash(other)
 
     @classmethod
-    def _indexed_fields(cls):
-        return sorted(field_name for field_name, field in cls.model_fields.items() if "indexed" in field.metadata)
+    def indexed_fields(cls):
+        indexed_fields = {}
+
+        # Handle regular fields
+        for fieldname, field in cls.model_fields.items():
+            if any(isinstance(m, str) and m.startswith("indexed") for m in field.metadata):
+                indexed_fields[fieldname] = field.metadata
+
+        # Handle computed fields
+        for fieldname, field in cls.model_computed_fields.items():
+            return_type = field.return_type
+            if get_origin(return_type) is Annotated:
+                type_args = get_args(return_type)
+                metadata = list(type_args[1:])  # Skip the first arg (the actual type)
+                if any(isinstance(m, str) and m.startswith("indexed") for m in metadata):
+                    indexed_fields[fieldname] = metadata
+
+        return indexed_fields
 
     # we keep these because they were a lot of work to make and maybe someday they'll be useful again
 
@@ -70,9 +81,6 @@ class Event(BBOTBaseModel):
     host: Annotated[Optional[str], "indexed"] = None
     port: Optional[int] = None
     netloc: Optional[str] = None
-    # we store the host in reverse to allow for instant subdomain queries
-    # this works because indexes are left-anchored, but we need to search starting from the right side
-    reverse_host: Annotated[Optional[str], "indexed"] = ""
     resolved_hosts: Optional[List] = None
     dns_children: Optional[dict] = None
     web_spider_distance: int = 10
@@ -98,6 +106,17 @@ def get_data(self):
     def __hash__(self):
         return hash(self.id)
 
+    @computed_field
+    @property
+    def reverse_host(self) -> Annotated[Optional[str], "indexed"]:
+        """
+        We store the host in reverse to allow for instant subdomain queries
+        This works because indexes are left-anchored, but we need to search starting from the right side
+        """
+        if self.host:
+            return self.host[::-1]
+        return None
+
 
 ### SCAN ###
 

bbot/modules/output/http.py

@@ -1,3 +1,5 @@
+from omegaconf import OmegaConf
+
 from bbot.models.pydantic import Event
 from bbot.modules.output.base import BaseOutputModule
 
@@ -32,7 +34,7 @@ async def setup(self):
         self.url = self.config.get("url", "")
         self.method = self.config.get("method", "POST")
         self.timeout = self.config.get("timeout", 10)
-        self.headers = self.config.get("headers", {})
+        self.headers = OmegaConf.to_object(self.config.get("headers", OmegaConf.create()))
         bearer = self.config.get("bearer", "")
         if bearer:
             self.headers["Authorization"] = f"Bearer {bearer}"

bbot/modules/output/mongo.py

@@ -5,6 +5,10 @@
 
 
 class Mongo(BaseOutputModule):
+    """
+    docker run --rm -p 27017:27017 mongo
+    """
+
     watched_events = ["*"]
     meta = {
         "description": "Output scan data to a MongoDB database",
@@ -48,11 +52,11 @@ async def setup(self):
         self.targets_collection = self.db[f"{self.collection_prefix}targets"]
 
         # Build an index for each field in reverse_host and host
-        for field_name, field in Event.model_fields.items():
-            if "indexed" in field.metadata:
-                unique = "unique" in field.metadata
-                await self.events_collection.create_index([(field_name, 1)], unique=unique)
-                self.verbose(f"Index created for field: {field_name} (unique={unique})")
+        for fieldname, metadata in Event.indexed_fields().items():
+            if "indexed" in metadata:
+                unique = "unique" in metadata
+                await self.events_collection.create_index([(fieldname, 1)], unique=unique)
+                self.verbose(f"Index created for field: {fieldname} (unique={unique})")
 
         return True
 
@@ -65,6 +69,7 @@ async def handle_event(self, event):
                 break
             except Exception as e:
                 self.warning(f"Error inserting event into MongoDB: {e}, retrying...")
+                self.trace()
                 await self.helpers.sleep(1)
 
         if event.type == "SCAN":

bbot/test/test_step_1/test_db_models.py

@@ -18,7 +18,7 @@ def test_pydantic_models(events, bbot_scanner):
     assert utc_now2.timestamp() == utc_now.timestamp()
 
     test_event = Event(**events.ipv4.json())
-    assert sorted(test_event._indexed_fields()) == [
+    assert sorted(test_event.indexed_fields()) == [
         "data",
         "host",
         "id",

bbot/test/test_step_2/module_tests/test_module_elastic.py

@@ -103,9 +103,9 @@ async def check(self, module_test, events):
                 # Events don't match exactly because the elastic ones have reverse_host and inserted_at
                 assert events_json != db_events_pydantic
                 for db_event in db_events_pydantic:
-                    db_event.pop("reverse_host")
-                    db_event.pop("inserted_at")
-                    db_event.pop("archived")
+                    db_event.pop("reverse_host", None)
+                    db_event.pop("inserted_at", None)
+                    db_event.pop("archived", None)
                 # They should match after removing reverse_host
                 assert events_json == db_events_pydantic, "Events do not match"
 

bbot/test/test_step_2/module_tests/test_module_mongo.py

@@ -76,7 +76,7 @@ async def check(self, module_test, events):
             # make sure the collection has all the right indexes
             cursor = events_collection.list_indexes()
             indexes = await cursor.to_list(length=None)
-            for field in Event._indexed_fields():
+            for field in Event.indexed_fields():
                 assert any(field in index["key"] for index in indexes), f"Index for {field} not found"
 
             ### EVENTS ###
@@ -116,9 +116,9 @@ async def check(self, module_test, events):
             # Events don't match exactly because the mongo ones have reverse_host and inserted_at
             assert events_json != db_events_pydantic
             for db_event in db_events_pydantic:
-                db_event.pop("reverse_host")
-                db_event.pop("inserted_at")
-                db_event.pop("archived")
+                db_event.pop("reverse_host", None)
+                db_event.pop("inserted_at", None)
+                db_event.pop("archived", None)
             # They should match after removing reverse_host
             assert events_json == db_events_pydantic, "Events do not match"