Merge pull request #27 from blacklanternsecurity/authentication

Add Authentication

Author: TheTechromancer

.github/workflows/docker-tests.yml

@@ -0,0 +1,28 @@
+name: docker tests
+on:
+  push:
+    branches:
+      - stable
+      - dev
+  pull_request:
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event_name }}-${{ github.event.pull_request.number || github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  test:
+    runs-on: self-hosted
+    steps:
+      - uses: actions/checkout@v4
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: "3.11"
+      - name: Install dependencies
+        run: |
+          pip install poetry
+          poetry install
+      - name: Run tests
+        run: |
+          BBOT_SERVER_TEST_DOCKER_COMPOSE=true poetry run pytest --disable-warnings --log-cli-level=INFO -k test_docker_compose

README.md

@@ -46,7 +46,53 @@ Note: this requires Docker and Docker Compose to be installed.
 bbctl server start
 ```
 
-## Start a scan (direct from BBOT CLI)
+## Interacting with BBOT Server Remotely (Multiplayer)
+
+By default, BBOT Server listens on localhost. Use `--listen` to expose it to the network:
+
+```bash
+bbctl server start --listen 0.0.0.0
+```
+
+### Authentication
+
+The first time you start BBOT Server, an API key will be auto generated and put into `~/.config/bbot_server/config.yml`:
+
+```yaml
+# ~/.config/bbot_server/config.yml
+
+# list of API keys to be considered valid
+api_keys:
+  - 4aa8b3c2-9b4d-4208-890c-4ce9ad3b4710
+```
+
+The `api_keys` value in `config.yml` is used by both the server (as a database of valid API keys), and by the client (it will pick one from the list and use it). Normally it just works and you don't have to mess with it. But to access BBOT Server remotely, you'll need to copy the API key from the server onto your local system, along with its URL:
+
+```yaml
+# ~/.config/bbot_server/config.yml
+url: http://1.2.3.4:8807/v1/
+api_keys:
+  - deadbeef-9b4d-4208-890c-4ce9ad3b4710
+```
+
+This tells `bbctl` (the client) where the server is, and gives it the means to authenticate.
+
+### Adding and Revoking API Keys
+
+API keys can be added and removed if you are on the server machine:
+
+```bash
+# add an API key
+bbctl server apikey add
+
+# list API keys
+bbctl server apikey list
+
+# revoke an API key
+bbctl server apikey delete deadbeef-9b4d-4208-890c-4ce9ad3b4710
+```
+
+## Send a BBOT Scan to the Server
 
 You can output a BBOT scan directly to BBOT server:
 
@@ -146,6 +192,23 @@ bbctl scan list
 bbctl scan cancel "demonic_jimmy"
 ```
 
+## Targets
+
+BBOT server categorizes its assets by target.
+
+You can list targets like so:
+
+```bash
+# List targets
+bbctl target list
+
+# Create a new target
+bbctl target create --seeds seeds.txt --blacklist blacklist.txt --name custom_target
+
+# List only the assets that match your new target
+bbctl asset list --target custom_target
+```
+
 ## Custom triggers
 
 You can kick off a custom command or bash script whenever a certain activity happens, such as when a new technology or open port is discovered.

bbot_server/agent/agent.py

@@ -8,17 +8,17 @@
 from typing import Callable, Any
 from urllib.parse import urlparse, urlunparse, urljoin
 
+import bbot_server.config as bbcfg
 from bbot.scanner import Scanner, Preset
 from bbot.scanner.dispatcher import Dispatcher
 from bbot.constants import get_scan_status_code, get_scan_status_name, SCAN_STATUS_NOT_STARTED, SCAN_STATUS_FAILED
 
-from bbot_server.config import BBOT_SERVER_CONFIG
 from bbot_server.errors import BBOTServerValueError
 from bbot_server.utils.async_utils import async_to_sync_class
 from bbot_server.models.agent_models import AgentResponse
 
-default_server_url = BBOT_SERVER_CONFIG.get("url", "http://localhost:8807/v1/")
-default_bbot_preset = BBOT_SERVER_CONFIG.get("agent", {}).get("base_preset", {})
+api_key = bbcfg.get_api_key()
+default_bbot_preset = bbcfg.BBOT_SERVER_CONFIG.get("agent", {}).get("base_preset", {})
 
 log = logging.getLogger("bbot_server.agent")
 
@@ -219,9 +219,7 @@ def make_agent_preset(self):
             output_modules=["http"],
             config={
                 "modules": {
-                    "http": {
-                        "url": self.scan_output_url,
-                    }
+                    "http": {"url": self.scan_output_url, "headers": {bbcfg.API_KEY_NAME: bbcfg.get_api_key()}}
                 }
             },
         )
@@ -234,7 +232,9 @@ async def loop(self):
             self.log.info(f"Agent {self.name} connecting to {self.websocket_dock_url}...")
             # "async for" will use websocket's builtin retry/reconnect mechanism, with exponential backoff
             # https://websockets.readthedocs.io/en/stable/reference/asyncio/client.html
-            async for websocket in websockets.connect(self.websocket_dock_url):
+            async for websocket in websockets.connect(
+                self.websocket_dock_url, additional_headers={bbcfg.API_KEY_NAME: api_key}
+            ):
                 self.log.info(f"Agent {self.name} successfully connected to {self.websocket_dock_url}")
                 self.websocket = websocket
                 try:

bbot_server/api/_fastapi.py

@@ -1,9 +1,14 @@
-from fastapi import FastAPI
+import logging
 from contextlib import asynccontextmanager
+from fastapi.openapi.utils import get_openapi
+from fastapi import FastAPI, HTTPException, Depends, Request, WebSocket
 from fastapi.responses import RedirectResponse, ORJSONResponse
 
+import bbot_server.config as bbcfg
 from bbot_server.errors import BBOTServerError, handle_bbot_server_error
 
+log = logging.getLogger("bbot_server.api.fastapi")
+
 # from fastapi.openapi.docs import get_swagger_ui_html, get_redoc_html
 
 app_kwargs = {
@@ -12,12 +17,44 @@
     "debug": True,
 }
 
+# API key header
+# api_key_header = APIKeyHeader(name=bbcfg.API_KEY_NAME, auto_error=False)
+
+
+def api_key_dependency_http(request: Request = None):
+    if request is not None:
+        api_key = request.headers.get(bbcfg.API_KEY_NAME, "")
+        if not api_key:
+            raise HTTPException(status_code=401, detail="API key is required")
+        valid, reason = bbcfg.check_api_key(api_key)
+        if not valid:
+            raise HTTPException(status_code=401, detail=reason)
+        return api_key
+
+
+async def api_key_dependency_websocket(websocket: WebSocket = None):
+    if websocket is not None:
+        api_key = websocket.headers.get(bbcfg.API_KEY_NAME, "")
+        if not api_key:
+            await websocket.close(code=3000, reason="API key is required")
+            raise HTTPException(status_code=401, detail="API key is required")
+        valid, reason = bbcfg.check_api_key(api_key)
+        if not valid:
+            await websocket.close(code=3000, reason=reason)
+            raise HTTPException(status_code=401, detail=reason)
+        return api_key
 
-def make_app(config=None):
+
+# Dependency to verify the API key
+# async def api_key_dependency(api_key: str = Security(api_key_header)):
+#     await verify_api_key(api_key)
+
+
+def make_app():
     from bbot_server.api.mcp import make_mcp_server
     from bbot_server.applets import BBOTServerRootApplet
 
-    app_root = BBOTServerRootApplet(config=config)
+    app_root = BBOTServerRootApplet()
     app_root._is_main_server = True
 
     @asynccontextmanager
@@ -31,9 +68,52 @@ async def lifespan(app: FastAPI):
         root_path="/v1",
         openapi_tags=app_root.tags_metadata,
         default_response_class=ORJSONResponse,
+        dependencies=[Depends(api_key_dependency_http), Depends(api_key_dependency_websocket)],
         **app_kwargs,
     )
 
+    # Customize OpenAPI to better document the authentication
+    def custom_openapi():
+        if app.openapi_schema:
+            return app.openapi_schema
+
+        openapi_schema = get_openapi(
+            title=app.title,
+            version=app.version,
+            description=app.description,
+            routes=app.routes,
+        )
+
+        # Prepend root_path to all paths in the OpenAPI schema
+        root_path = app.root_path or ""
+        if root_path and root_path != "/":
+            new_paths = {}
+            for path, path_item in openapi_schema["paths"].items():
+                if not path.startswith(root_path):
+                    new_path = root_path + path
+                else:
+                    new_path = path
+                new_paths[new_path] = path_item
+            openapi_schema["paths"] = new_paths
+
+        # Add security scheme to OpenAPI schema
+        openapi_schema["components"]["securitySchemes"] = {
+            "APIKeyHeader": {
+                "type": "apiKey",
+                "in": "header",
+                "name": bbcfg.API_KEY_NAME,
+                "description": "API key authentication",
+            }
+        }
+
+        # Add global security requirement
+        openapi_schema["security"] = [{"APIKeyHeader": []}]
+
+        app.openapi_schema = openapi_schema
+        return app.openapi_schema
+
+    app.openapi = custom_openapi
+
     app.include_router(app_root.router)
 
     # add MCP server to the fastapi app
@@ -67,8 +147,8 @@ async def docs_redirect():
     return app, lifespan
 
 
-def make_server_app(config=None):
-    app, lifespan = make_app(config=config)
+def make_server_app():
+    app, lifespan = make_app()
 
     # includes the /v1 prefix
     server_app = FastAPI(

bbot_server/applets/_base.py

@@ -449,6 +449,7 @@ def _add_custom_routes(self):
             # see if the value has an "_endpoint" attribute
             endpoint = getattr(function, "_endpoint", None)
             # if it's a callable function and it has _endpoint, it's an @api_endpoint
+
             if endpoint is not None:
                 fastapi_kwargs = dict(getattr(function, "_kwargs", {}))
                 endpoint_type = fastapi_kwargs.pop("type", "http")

bbot_server/applets/_root.py

@@ -1,14 +1,13 @@
-from omegaconf import OmegaConf
+import bbot_server.config as bbcfg
 
 from bbot_server.applets._base import BaseApplet
-from bbot_server.config import BBOT_SERVER_CONFIG
 
-# assets imports
+# applet imports
+from bbot_server.applets.stats import StatsApplet
 from bbot_server.applets.assets import AssetsApplet
 from bbot_server.applets.events import EventsApplet
 from bbot_server.applets.scans import ScansApplet
 from bbot_server.applets.activity import ActivityApplet
-from bbot_server.applets.stats import StatsApplet
 
 
 class RootApplet(BaseApplet):
@@ -24,11 +23,9 @@ def __init__(self, config=None, **kwargs):
         """
         "config" can be either a dictionary or an omegaconf object
         """
-        super().__init__(**kwargs)
         if config is not None:
-            self._config = OmegaConf.merge(BBOT_SERVER_CONFIG, config)
-        else:
-            self._config = BBOT_SERVER_CONFIG
+            bbcfg.refresh_config(config)
+        super().__init__(**kwargs)
         self._interface_type = "python"
         self._mcp = None
 
@@ -41,26 +38,26 @@ async def setup(self):
                 from bbot_server.store.user_store import UserStore
                 from bbot_server.store.asset_store import AssetStore
 
-                self.asset_store = AssetStore(self.global_config)
+                self.asset_store = AssetStore()
                 await self.asset_store.setup()
                 self.asset_db = self.asset_store.db
                 self.asset_fs = self.asset_store.fs
 
-                self.user_store = UserStore(self.global_config)
+                self.user_store = UserStore()
                 await self.user_store.setup()
                 self.user_db = self.user_store.db
                 self.user_fs = self.user_store.fs
 
             # set up event store
             from bbot_server.event_store import EventStore
 
-            self.event_store = EventStore(self.global_config)
+            self.event_store = EventStore()
             await self.event_store.setup()
 
             # set up NATS client
             from bbot_server.message_queue import MessageQueue
 
-            self.message_queue = MessageQueue(self.global_config)
+            self.message_queue = MessageQueue()
             await self.message_queue.setup()
 
         await self._setup()
@@ -70,6 +67,10 @@ async def setup(self):
     def config(self):
         return self._config
 
+    @property
+    def _config(self):
+        return bbcfg.BBOT_SERVER_CONFIG
+
     async def cleanup(self):
         if self.is_native:
             await self.asset_store.cleanup()