From d5d5f5a70b86b562e110fba9ad1e5aeff31a9209 Mon Sep 17 00:00:00 2001 From: Salad Dais Date: Sat, 1 May 2021 00:03:23 +0000 Subject: [PATCH] Improve command handling testing --- hippolyzer/lib/proxy/addons.py | 16 +++++++++++++--- .../lib/proxy/tests/integration/__init__.py | 1 + .../lib/proxy/tests/integration/test_addons.py | 17 +++++++++++++++-- 3 files changed, 29 insertions(+), 5 deletions(-) diff --git a/hippolyzer/lib/proxy/addons.py b/hippolyzer/lib/proxy/addons.py index cd27993..a4b65b9 100644 --- a/hippolyzer/lib/proxy/addons.py +++ b/hippolyzer/lib/proxy/addons.py @@ -73,9 +73,11 @@ class AddonManager: _SUBPROCESS: bool = False _REPL_TASK: Optional[asyncio.Task] = None _HOT_RELOADING_STACK: Set[str] = set() + _SWALLOW_ADDON_EXCEPTIONS: bool = True @classmethod - def init(cls, addon_script_paths, session_manager, addon_objects=None, subprocess=False): + def init(cls, addon_script_paths, session_manager, addon_objects=None, subprocess=False, + swallow_addon_exceptions=True): addon_objects = addon_objects or [] cls.BASE_ADDON_SPECS.clear() cls.FRESH_ADDON_MODULES.clear() @@ -85,6 +87,7 @@ class AddonManager: # Are we a child process from `multiprocessing`? cls._SUBPROCESS = subprocess cls._HOT_RELOADING_STACK.clear() + cls._SWALLOW_ADDON_EXCEPTIONS = swallow_addon_exceptions for path in addon_script_paths: # defer reloading until we've collected them all cls.load_addon_from_path(path, reload=False) @@ -369,6 +372,8 @@ class AddonManager: return hook_func(*args, **kwargs) except: logging.exception("Exploded in %r's %s hook" % (addon, hook_name)) + if not cls._SWALLOW_ADDON_EXCEPTIONS: + raise @classmethod def _show_message(cls, text: str, session: Optional[Session] = None): @@ -377,7 +382,8 @@ class AddonManager: try: show_message(text, session=session) except: - pass + if not cls._SWALLOW_ADDON_EXCEPTIONS: + raise @classmethod def _show_error(cls, text: str, session: Optional[Session] = None): @@ -389,13 +395,15 @@ class AddonManager: cls._reload_addons() if message.name == "ChatFromViewer" and "ChatData" in message: if message["ChatData"]["Channel"] == cls.COMMAND_CHANNEL: + region.circuit.drop_message(message) with addon_ctx.push(session, region): try: cls._handle_command(session, region, message["ChatData"]["Message"]) except Exception as e: LOG.exception(f'Failed while handling command {message["ChatData"]["Message"]}') cls._show_message(str(e), session) - region.circuit.drop_message(message) + if not cls._SWALLOW_ADDON_EXCEPTIONS: + raise return True if message.name == "ChatFromSimulator" and "ChatData" in message: chat: str = message["ChatData"]["Message"] @@ -416,6 +424,8 @@ class AddonManager: return True except: LOG.exception(f"Failed while handling command {chat!r}") + if not cls._SWALLOW_ADDON_EXCEPTIONS: + raise with addon_ctx.push(session, region): return cls._call_all_addon_hooks("handle_lludp_message", session, region, message) diff --git a/hippolyzer/lib/proxy/tests/integration/__init__.py b/hippolyzer/lib/proxy/tests/integration/__init__.py index 3a2b432..35c7241 100644 --- a/hippolyzer/lib/proxy/tests/integration/__init__.py +++ b/hippolyzer/lib/proxy/tests/integration/__init__.py @@ -51,6 +51,7 @@ class BaseIntegrationTest(unittest.IsolatedAsyncioTestCase): self.session_manager.claim_session(self.session.id) self.session.open_circuit(self.client_addr, self.region_addr, self.protocol.transport) + self.session.main_region = self.session.regions[-1] def _msg_to_datagram(self, msg: ProxiedMessage, src, dst, direction, socks_header=True): serialized = self.serializer.serialize(msg) diff --git a/hippolyzer/lib/proxy/tests/integration/test_addons.py b/hippolyzer/lib/proxy/tests/integration/test_addons.py index ee6b575..aa4ce22 100644 --- a/hippolyzer/lib/proxy/tests/integration/test_addons.py +++ b/hippolyzer/lib/proxy/tests/integration/test_addons.py @@ -2,7 +2,12 @@ from __future__ import annotations from hippolyzer.lib.base.message.message import Block from hippolyzer.lib.proxy import addon_ctx -from hippolyzer.lib.proxy.addon_utils import BaseAddon, SessionProperty +from hippolyzer.lib.proxy.addon_utils import ( + BaseAddon, + SessionProperty, + send_chat, + show_message, +) from hippolyzer.lib.proxy.addons import AddonManager from hippolyzer.lib.proxy.commands import handle_command from hippolyzer.lib.proxy.message import ProxiedMessage @@ -20,13 +25,18 @@ class MockAddon(BaseAddon): async def foobar(self, _session: Session, _region: ProxiedRegion, bar: str): self.bazquux = bar self.another = bar + send_chat(bar) + show_message(bar) class AddonIntegrationTests(BaseIntegrationTest): def setUp(self) -> None: super().setUp() self.addon = MockAddon() - AddonManager.init([], self.session_manager, [self.addon]) + AddonManager.init([], self.session_manager, [self.addon], swallow_addon_exceptions=False) + + def tearDown(self) -> None: + AddonManager.shutdown() def _fake_command(self, command: str) -> None: msg = ProxiedMessage( @@ -63,3 +73,6 @@ class AddonIntegrationTests(BaseIntegrationTest): _something = self.addon.bazquux # This has a default self.assertEqual(self.addon.another, "default") + # Should have sent out the two injected packets for inbound and outbound chat + # But not the original chatfromviewer from our command. + self.assertEqual(len(self.transport.packets), 2)