Merge pull request #761 from blacklanternsecurity/agent-fixes

Misc Agent Fixes
blacklanternsecurity · Sep 25, 2023 · 1c27596 · 1c27596
2 parents aaa3aba + 20edb21
commit 1c27596
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 7 deletions.
diff --git a/bbot/agent/agent.py b/bbot/agent/agent.py
@@ -51,7 +51,14 @@ async def ws(self, rebuild=False):
                 verbs = ("Rebuilding", "Rebuilt")
             url = f"{self.url}/control/"
             log.debug(f"{verbs[0]} websocket connection to {url}")
-            self._ws = await websockets.connect(url, **kwargs)
+            while 1:
+                try:
+                    self._ws = await websockets.connect(url, **kwargs)
+                    break
+                except Exception as e:
+                    log.error(f'Failed to establish websockets connection to URL "{url}": {e}')
+                    log.trace(traceback.format_exc())
+                    await asyncio.sleep(1)
             log.debug(f"{verbs[1]} websocket connection to {url}")
         return self._ws
 
@@ -70,6 +77,7 @@ async def start(self):
                     if message.command == "ping":
                         if self.scan is None:
                             await self.send({"conversation": str(message.conversation), "message_type": "pong"})
+                            continue
 
                     command_type = getattr(messages, message.command, None)
                     if command_type is None:
@@ -115,11 +123,7 @@ 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}/control/scan/{scan_id}/", "token": self.token}
-                        }
-                    }
+                    {"output_modules": {"websocket": {"url": f"{self.url}/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_cli.py b/bbot/test/test_step_1/test_cli.py
@@ -53,7 +53,10 @@ async def test_cli(monkeypatch, bbot_config):
     task = asyncio.create_task(cli._main())
     await asyncio.sleep(2)
     task.cancel()
-    await task
+    try:
+        await task
+    except asyncio.CancelledError:
+        pass
 
     # no args
     monkeypatch.setattr("sys.argv", ["bbot"])