Merge pull request #91 from blacklanternsecurity/config-rework-2

Config Rework

Author: TheTechromancer

README.md

@@ -379,7 +379,7 @@ if __name__ == "__main__":
 When running tests, first start MongoDB and Redis via Docker:
 
 ```bash
-docker run --rm -p 27017:27017 mongo
+docker run --ulimit nofile=64000:64000 --rm -p 27017:27017 mongo
 docker run --rm -p 6379:6379 redis
 ```
 

bbot_server/api/_fastapi.py

@@ -5,7 +5,7 @@
 from fastapi import FastAPI, HTTPException, Depends, Request, WebSocket
 
 from bbot_server import modules  # noqa: F401
-import bbot_server.config as bbcfg
+from bbot_server.config import BBOT_SERVER_CONFIG as bbcfg
 from bbot_server.errors import BBOTServerError, handle_bbot_server_error
 
 log = logging.getLogger("bbot_server.api.fastapi")
@@ -19,12 +19,14 @@
 }
 
 # API key header
-# api_key_header = APIKeyHeader(name=bbcfg.API_KEY_NAME, auto_error=False)
+# api_key_header = APIKeyHeader(name=bbcfg.auth_header, auto_error=False)
 
 
 def api_key_dependency_http(request: Request = None):
+    if not bbcfg.auth_enabled:
+        return
     if request is not None:
-        api_key = request.headers.get(bbcfg.API_KEY_NAME, "")
+        api_key = request.headers.get(bbcfg.auth_header, "")
         if not api_key:
             raise HTTPException(status_code=401, detail="API key is required")
         valid, reason = bbcfg.check_api_key(api_key)
@@ -34,15 +36,15 @@ def api_key_dependency_http(request: Request = None):
 
 
 async def api_key_dependency_websocket(websocket: WebSocket = None):
+    if not bbcfg.auth_enabled:
+        return
     if websocket is not None:
-        api_key = websocket.headers.get(bbcfg.API_KEY_NAME, "")
+        api_key = websocket.headers.get(bbcfg.auth_header, "")
         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
 
 
@@ -102,7 +104,7 @@ def custom_openapi():
             "APIKeyHeader": {
                 "type": "apiKey",
                 "in": "header",
-                "name": bbcfg.API_KEY_NAME,
+                "name": bbcfg.auth_header,
                 "description": "API key authentication",
             }
         }

bbot_server/applets/_root.py

@@ -1,5 +1,4 @@
-import bbot_server.config as bbcfg
-
+from bbot_server.config import BBOT_SERVER_CONFIG as bbcfg
 from bbot_server.applets.base import BaseApplet
 
 
@@ -12,10 +11,10 @@ class RootApplet(BaseApplet):
 
     def __init__(self, config=None, **kwargs):
         """
-        "config" can be either a dictionary or an omegaconf object
+        "config" can be a dictionary of config overrides
         """
         if config is not None:
-            bbcfg.refresh_config(config)
+            bbcfg.refresh(**config)
         super().__init__(**kwargs)
         self._interface_type = "python"
         self._mcp = None
@@ -52,7 +51,7 @@ def config(self):
 
     @property
     def _config(self):
-        return bbcfg.BBOT_SERVER_CONFIG
+        return bbcfg
 
     async def cleanup(self):
         if self.is_native:

bbot_server/applets/_routing.py

@@ -7,7 +7,7 @@
 from contextlib import suppress
 from fastapi.responses import StreamingResponse
 from starlette.websockets import WebSocketDisconnect
-import bbot_server.config as bbcfg
+from bbot_server.config import BBOT_SERVER_CONFIG as bbcfg
 from bbot_server.api.mcp import MCP_ENDPOINTS
 from bbot_server.utils.misc import smart_encode
 from bbot_server.errors import BBOTServerValueError
@@ -230,7 +230,7 @@ def wrapped_function(self):
         @functools.wraps(self.orig_function)
         async def websocket_auth_wrapper(websocket: WebSocket, *args, **kwargs):
             await websocket.accept()
-            api_key = websocket.headers.get(bbcfg.API_KEY_NAME, "")
+            api_key = websocket.headers.get(bbcfg.auth_header, "")
             valid, reason = bbcfg.check_api_key(api_key)
             if valid:
                 await self.orig_function(websocket, *args, **kwargs)
@@ -260,7 +260,7 @@ async def websocket_wrapper(websocket: WebSocket, *args, **kwargs):
             """
             try:
                 await websocket.accept()
-                api_key = websocket.headers.get(bbcfg.API_KEY_NAME, "")
+                api_key = websocket.headers.get(bbcfg.auth_header, "")
                 valid, reason = bbcfg.check_api_key(api_key)
                 if not valid:
                     await websocket.close(code=3000, reason=reason)
@@ -305,7 +305,7 @@ def wrapped_function(self):
     async def websocket_wrapper(self, websocket: WebSocket):
         try:
             await websocket.accept()
-            api_key = websocket.headers.get(bbcfg.API_KEY_NAME, "")
+            api_key = websocket.headers.get(bbcfg.auth_header, "")
             valid, reason = bbcfg.check_api_key(api_key)
             if not valid:
                 await websocket.close(code=3000, reason=reason)

bbot_server/applets/base.py

@@ -658,7 +658,7 @@ def global_config(self):
 
     @property
     def config(self):
-        return self.global_config.get("modules", {}).get(self.name, {})
+        return self.global_config.modules.get(self.name, {})
 
     @property
     def tag(self):

bbot_server/cli/base.py

@@ -7,7 +7,7 @@
 from typing import Annotated  # noqa
 from functools import cached_property, wraps
 
-import bbot_server.config as bbcfg
+from bbot_server.config import BBOT_SERVER_CONFIG as bbcfg
 from bbot_server.utils.misc import timestamp_to_human, seconds_to_human
 
 
@@ -168,7 +168,7 @@ def highlight_yaml(self, data, **kwargs):
         Highlight a YAML string with rich
         """
         if not isinstance(data, str):
-            data = yaml.dump(data, indent=2)
+            data = yaml.safe_dump(data, indent=2, sort_keys=False)
         if not "background_color" in kwargs:
             kwargs["background_color"] = "default"
         if not "theme" in kwargs:

bbot_server/cli/bbctl.py

@@ -3,12 +3,9 @@
 import asyncio
 import logging
 import traceback
-from pathlib import Path
-from omegaconf import OmegaConf
 from rich.console import Console
 from functools import cached_property
 
-import bbot_server.config as bbcfg
 from bbot_server.cli.base import BaseBBCTL, Annotated, Option
 from bbot_server.errors import BBOTServerError, BBOTServerUnauthorizedError
 
@@ -40,45 +37,33 @@ def main(
         self.silent = silent
         self.color = color
         self.debug = debug
-        self.config_path = None
-        # command line arg takes precedence over environment variable
-        custom_config = config or os.environ.get("BBOT_SERVER_CONFIG", "")
-        if custom_config:
-            try:
-                self.config_path = Path(custom_config)
-                self.bbcfg.update_config_path(self.config_path)
-            except Exception as e:
-                raise BBOTServerError(f"Error loading config file at {self.config_path}: {e}")
-        else:
-            self.config_path = bbcfg.BBOT_SERVER_CONFIG_PATH
+        if config:
+            os.environ["BBOT_SERVER_CONFIG"] = str(config)
+            self.bbcfg.refresh()
         if self.debug:
             logging.getLogger().setLevel(logging.DEBUG)
-        if server_url is not None and server_url != bbcfg.BBOT_SERVER_URL:
+        if server_url is not None and server_url != self.bbcfg.url:
             self._config.url = server_url
         self.server_url = self.config.url
 
         self._stdout = Console(file=sys.stdout, highlight=False, color_system=("auto" if self.color else None))
         self._stderr = Console(file=sys.stderr, highlight=False, color_system=("auto" if self.color else None))
 
         if current_config:
-            self.print_yaml(OmegaConf.to_yaml(self.config))
+            self.print_yaml(self.config.model_dump())
             return
 
     @cached_property
     def bbot_server(self):
-        bbot_server_kwargs = {}
-        if self.config:
-            bbot_server_kwargs["config"] = self.config
-
         from bbot_server import BBOTServer
 
-        bbot_server = BBOTServer(interface="http", url=self.server_url, synchronous=True, **bbot_server_kwargs)
+        bbot_server = BBOTServer(interface="http", url=self.server_url, synchronous=True)
         bbot_server.setup()
         return bbot_server
 
     @property
     def _config(self):
-        return self.bbcfg.BBOT_SERVER_CONFIG
+        return self.bbcfg
 
 
 log = logging.getLogger("bbot_server.bbctl")

bbot_server/cli/themes.py

@@ -1,20 +1,20 @@
 import typer
-from textual.theme import Theme
+# from textual.theme import Theme
 
 COLOR = "bold dark_orange"
 DARK_COLOR = "grey50"
 
-# textual theme
-TEXTUAL_THEME = Theme(
-    name="bbot",
-    primary="#000000",
-    secondary="#1a1a1a",
-    accent="#FF8400",
-    warning="#FF8400",
-    error="#ff4500",
-    success="#FF8400",
-    foreground="#ffffff",
-)
+# # textual theme
+# TEXTUAL_THEME = Theme(
+#     name="bbot",
+#     primary="#000000",
+#     secondary="#1a1a1a",
+#     accent="#FF8400",
+#     warning="#FF8400",
+#     error="#ff4500",
+#     success="#FF8400",
+#     foreground="#ffffff",
+# )
 
 # typer theme
 typer.rich_utils.STYLE_OPTION = "bold dark_orange"