From 7ea2311b175bc3856c586284e8cf6c3830a9b27b Mon Sep 17 00:00:00 2001
From: TheTechromancer <thetechromancer@protonmail.com>
Date: Mon, 25 Sep 2023 12:04:19 -0400
Subject: [PATCH 1/4] fixed unhandled websockets error in agent

---
 bbot/agent/agent.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/bbot/agent/agent.py b/bbot/agent/agent.py
index f30fb94d15..09eff2a03b 100644
--- 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
 

From 69189788b372876d14a46fedd4d6ba9fd668d966 Mon Sep 17 00:00:00 2001
From: TheTechromancer <thetechromancer@protonmail.com>
Date: Mon, 25 Sep 2023 12:08:32 -0400
Subject: [PATCH 2/4] fix 'unknown command: ping' error

---
 bbot/agent/agent.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/bbot/agent/agent.py b/bbot/agent/agent.py
index 09eff2a03b..819388274c 100644
--- a/bbot/agent/agent.py
+++ b/bbot/agent/agent.py
@@ -77,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:

From 1d92059d615ef33b2c3fbe531b42b3fba301890f Mon Sep 17 00:00:00 2001
From: TheTechromancer <thetechromancer@protonmail.com>
Date: Mon, 25 Sep 2023 12:18:12 -0400
Subject: [PATCH 3/4] fix websocket output URL from agent

---
 bbot/agent/agent.py | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/bbot/agent/agent.py b/bbot/agent/agent.py
index 819388274c..1c8debc1e0 100644
--- a/bbot/agent/agent.py
+++ b/bbot/agent/agent.py
@@ -123,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)

From 20edb218eb9dbe93c6840b9d168e06b3c4e5375a Mon Sep 17 00:00:00 2001
From: TheTechromancer <thetechromancer@protonmail.com>
Date: Mon, 25 Sep 2023 12:41:06 -0400
Subject: [PATCH 4/4] fix agent tests

---
 bbot/test/test_step_1/test_cli.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/bbot/test/test_step_1/test_cli.py b/bbot/test/test_step_1/test_cli.py
index 31b681c177..0ccc94887e 100644
--- 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"])