Merge pull request #682 from blacklanternsecurity/681-agent-url-for-w…

…ebsocket-fix add '/control/' to websocket url
blacklanternsecurity · Sep 1, 2023 · 0da5a94 · 0da5a94
2 parents 2ad3555 + 7bb5597
commit 0da5a94
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 5 deletions.
diff --git a/bbot/agent/agent.py b/bbot/agent/agent.py
@@ -49,9 +49,10 @@ async def ws(self, rebuild=False):
             verbs = ("Building", "Built")
             if rebuild:
                 verbs = ("Rebuilding", "Rebuilt")
-            log.debug(f"{verbs[0]} websocket connection to {self.url}")
-            self._ws = await websockets.connect(self.url, **kwargs)
-            log.debug(f"{verbs[1]} websocket connection to {self.url}")
+            url = f"{self.url}/control/"
+            log.debug(f"{verbs[0]} websocket connection to {url}")
+            self._ws = await websockets.connect(url, **kwargs)
+            log.debug(f"{verbs[1]} websocket connection to {url}")
         return self._ws
 
     async def start(self):
@@ -114,7 +115,11 @@ async def start_scan(self, scan_id, name=None, targets=[], modules=[], output_mo
                     f"Starting scan with targets={targets}, modules={modules}, output_modules={output_modules}"
                 )
                 output_module_config = OmegaConf.create(
-                    {"output_modules": {"websocket": {"url": f"{self.url}/scan/{scan_id}/", "token": self.token}}}
+                    {
+                        "output_modules": {
+                            "websocket": {"url": f"{self.url}/control/scan/{scan_id}/", "token": self.token}
+                        }
+                    }
                 )
                 config = OmegaConf.create(config)
                 config = OmegaConf.merge(self.config, config, output_module_config)

diff --git a/bbot/test/test_step_1/test_agent.py b/bbot/test/test_step_1/test_agent.py
@@ -20,7 +20,7 @@ async def websocket_handler(websocket, path, scan_done=None):
     # control channel or event channel?
     control = True
 
-    if path == "/" and first_run:
+    if path == "/control/" and first_run:
         # test ping
         await websocket.send(json.dumps({"conversation": "90196cc1-299f-4555-82a0-bc22a4247590", "command": "ping"}))
         _first_run = False