diff --git a/bbot/agent/agent.py b/bbot/agent/agent.py index f30fb94d1..1c8debc1e 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 @@ -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 index 31b681c17..0ccc94887 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"])