diff --git a/bbot/models/pydantic.py b/bbot/models/pydantic.py index 4ec7a3dc65..68e71d493e 100644 --- a/bbot/models/pydantic.py +++ b/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 ### diff --git a/bbot/modules/output/http.py b/bbot/modules/output/http.py index 3b61556cb6..28fa917fc7 100644 --- a/bbot/modules/output/http.py +++ b/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}" diff --git a/bbot/modules/output/mongo.py b/bbot/modules/output/mongo.py index 118ca82378..f90c4aad53 100644 --- a/bbot/modules/output/mongo.py +++ b/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": diff --git a/bbot/test/test_step_1/test_db_models.py b/bbot/test/test_step_1/test_db_models.py index 0e175f3418..d453fa81e1 100644 --- a/bbot/test/test_step_1/test_db_models.py +++ b/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", diff --git a/bbot/test/test_step_2/module_tests/test_module_elastic.py b/bbot/test/test_step_2/module_tests/test_module_elastic.py index b3af6e694a..9846f02246 100644 --- a/bbot/test/test_step_2/module_tests/test_module_elastic.py +++ b/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" diff --git a/bbot/test/test_step_2/module_tests/test_module_mongo.py b/bbot/test/test_step_2/module_tests/test_module_mongo.py index 978d79c4b3..25c317e57c 100644 --- a/bbot/test/test_step_2/module_tests/test_module_mongo.py +++ b/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"