Merge pull request #55 from blacklanternsecurity/tech-improvements

Filtering Improvements

Author: TheTechromancer

README.md

@@ -325,6 +325,71 @@ The SSE server listens at `http://localhost:8807/v1/mcp/`
 
 After connecting your AI client to BBOT Server, you can ask it sensible questions like, "Use MCP to get all the bbot findings", "what are the top open ports?", "what else can you do with BBOT MCP?", etc.
 
+## As a Python Library
+
+You can interact fully with BBOT Server as a Python library. It supports either local or remote connections, and the interface to both is identical:
+
+### Asynchronous
+
+```python
+import asyncio
+from bbot_server import BBOTServer
+
+async def main():
+    # talk directly to local MongoDB + Redis
+    bbot_server = BBOTServer(interface="python")
+
+    # or to a remote BBOT Server instance (config must contain a valid API key)
+    bbot_server = BBOTServer(interface="http", url="http://bbot:8807/v1/")
+
+    # one-time setup
+    await bbot_server.setup()
+
+    hosts = await bbot_server.get_hosts()
+    print(f"hosts: {hosts}")
+
+if __name__ == "__main__":
+    asyncio.run(main())
+```
+
+### Synchronous
+
+```python
+from bbot_server import BBOTServer
+
+if __name__ == "__main__":
+    # talk directly to local MongoDB + Redis
+    bbot_server = BBOTServer(interface="python", synchronous=True)
+
+    # or to a remote BBOT Server instance (config must contain a valid API key)
+    bbot_server = BBOTServer(interface="http", url="http://bbot:8807/v1/", synchronous=True)
+
+    # one-time setup
+    bbot_server.setup()
+
+    hosts = bbot_server.get_hosts()
+    print(f"hosts: {hosts}")
+```
+
+## Running Tests
+
+When running tests, first start MongoDB and Redis via Docker:
+
+```bash
+docker run --rm -p 27017:27017 mongo
+docker run --rm -p 6379:6379 redis
+```
+
+Then execute `pytest`:
+
+```bash
+# run all tests
+poetry run pytest -v
+
+# run specific tests
+poetry run pytest -v -k test_applet_scans
+```
+
 ## Screenshots
 
 *Tailing activities in real time*

bbot_server/applets/_root.py

@@ -6,8 +6,6 @@
 class RootApplet(BaseApplet):
     name = "Root Applet"
 
-    attach_to = ""
-
     _nested = False
 
     _route_prefix = ""

bbot_server/applets/base.py

@@ -79,6 +79,7 @@ class BaseApplet:
     model = None
 
     # which other applet should include this one
+    # leave blank to attach to the root applet
     attach_to = ""
 
     # whether to nest this applet under its parent

bbot_server/db/base.py

@@ -33,7 +33,10 @@ def db_config(self):
 
     @property
     def uri(self):
-        return self.db_config.get("uri", "")
+        uri = self.db_config.get("uri", "")
+        if not uri:
+            raise BBOTServerValueError(f"Database URI is missing from config: {self.db_config}")
+        return uri
 
     @property
     def db_name(self):

bbot_server/event_store/_base.py

@@ -7,9 +7,8 @@ class BaseEventStore(BaseDB):
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self.event_store_config = self.config.get("event_store", {})
-        self.archive_after_days = self.event_store_config.get("archive_after", 90)
-        self.archive_cron = self.event_store_config.get("archive_cron", "0 0 * * *")
+        self.archive_after_days = self.db_config.get("archive_after", 90)
+        self.archive_cron = self.db_config.get("archive_cron", "0 0 * * *")
 
     async def insert_event(self, event):
         if not isinstance(event, Event):
@@ -20,9 +19,26 @@ async def get_event(self, uuid: str):
         event = await self._get_event(uuid)
         return Event(**event)
 
-    async def get_events(self, host: str = None, type=None, min_timestamp=None, archived=False, active=True):
+    async def get_events(
+        self,
+        host: str = None,
+        domain: str = None,
+        type: str = None,
+        scan: str = None,
+        min_timestamp: float = None,
+        max_timestamp: float = None,
+        archived: bool = False,
+        active: bool = True,
+    ):
         async for event in self._get_events(
-            host=host, type=type, min_timestamp=min_timestamp, archived=archived, active=active
+            host=host,
+            domain=domain,
+            type=type,
+            scan=scan,
+            min_timestamp=min_timestamp,
+            max_timestamp=max_timestamp,
+            archived=archived,
+            active=active,
         ):
             yield Event(**event)
 

bbot_server/event_store/mongo.py

@@ -1,13 +1,13 @@
 from pymongo import WriteConcern, AsyncMongoClient
 
 
-from bbot_server.errors import BBOTServerNotFoundError
+from bbot_server.errors import BBOTServerNotFoundError, BBOTServerValueError
 from bbot_server.event_store._base import BaseEventStore
 
 
 class MongoEventStore(BaseEventStore):
     """
-    docker run --rm -p 27017:27017 mongo
+    docker run --rm -p 127.0.0.1:27017:27017 mongo
     """
 
     async def _setup(self):
@@ -28,21 +28,43 @@ async def _insert_event(self, event):
         event_json = event.model_dump()
         await self.collection.insert_one(event_json)
 
-    async def _get_events(self, host: str, type: str, min_timestamp: float, archived: bool, active: bool):
+    async def _get_events(
+        self,
+        host: str,
+        domain: str,
+        type: str,
+        scan: str,
+        min_timestamp: float,
+        max_timestamp: float,
+        active: bool,
+        archived: bool,
+    ):
         """
-        Get all events from the database, or if min_timestamp is provided, get the newest events up to that timestamp
+        Get all events from the database based on the provided filters
         """
         query = {}
         if type is not None:
-            query["type"] = {"$eq": type}
+            query["type"] = type
         if min_timestamp is not None:
             query["timestamp"] = {"$gte": min_timestamp}
+        # if both active and archived are true, we don't need to filter anything
         if not (active and archived):
+            # if both are false, we need to raise an error
             if not (active or archived):
-                raise ValueError("Must query at least one of active or archived")
+                raise BBOTServerValueError("Must query at least one of active or archived")
+            # otherwise if only one is true, we need to filter by the other
             query["archived"] = {"$eq": archived}
+        if max_timestamp is not None:
+            query["timestamp"] = {"$lte": max_timestamp}
+        if scan is not None:
+            query["scan"] = scan
         if host is not None:
             query["host"] = host
+        if domain is not None:
+            # match reverse_host with regex
+            reversed_host = domain[::-1]
+            query["reverse_host"] = {"$regex": f"^{reversed_host}(\\.|$)"}
+        self.log.debug(f"Querying events: {query}")
         async for event in self.collection.find(query):
             yield event
 

bbot_server/message_queue/redis.py

@@ -29,7 +29,7 @@ class RedisMessageQueue(BaseMessageQueue):
     - bbot:stream:{subject}: for persistent, tailable streams - e.g. events, activities
     - bbot:work:{subject}: for one-time messages, e.g. tasks
 
-    docker run --rm -p 6379:6379 redis
+    docker run --rm -p 127.0.0.1:6379:6379 redis
     """
 
     def __init__(self, *args, **kwargs):

bbot_server/modules/__init__.py

@@ -133,7 +133,7 @@ def load_python_file(file, namespace, module_dict, base_class_name, module_key_a
                         log.error(f"Module {value.__name__} does not define required attribute{module_key_attr}")
                     parent_name = getattr(value, "attach_to", "")
                     if not parent_name:
-                        log.error(f"Module {value.__name__} does not define required attribute 'attach_to'")
+                        parent_name = "root_applet"
                     module_family = module_dict.get(parent_name, {})
                     # if we get a duplicate module name, raise an error
                     if module_name in module_family:

bbot_server/modules/activity/activity_api.py

@@ -8,8 +8,7 @@
 class ActivityApplet(BaseApplet):
     name = "Activity"
     watched_activities = ["*"]
-    description = "Query and tail BBOT server activity"
-    attach_to = "root_applet"
+    description = "Query BBOT server activities"
     model = Activity
 
     async def handle_activity(self, activity: Activity, asset: Asset = None):

bbot_server/modules/agents/agents_api.py

@@ -252,7 +252,7 @@ async def _kickoff_queued_scans(self):
         return 0
 
     async def _kickoff_queued_scans_loop(self):
-        for i in range(1000):
+        while True:
             online_agents = await self.get_online_agents()
             online_agents = [str(agent.id) for agent in online_agents]
             if not online_agents: