Add hooks to allow swapping out transports

addon_examples/serialization_sanity_checker.py

@@ -29,10 +29,11 @@ class SerializationSanityChecker(BaseAddon):
         self.deserializer = UDPMessageDeserializer()
 
     def handle_proxied_packet(self, session_manager: SessionManager, packet: UDPPacket,
-                              session: Optional[Session], region: Optional[ProxiedRegion],
-                              message: Optional[Message]):
+                              session: Optional[Session], region: Optional[ProxiedRegion]):
         # Well this doesn't even parse as a message, can't do anything about it.
-        if message is None:
+        try:
+            message = self.deserializer.deserialize(packet.data)
+        except:
             LOG.error(f"Received unparseable message from {packet.src_addr!r}: {packet.data!r}")
             return
         try:

hippolyzer/apps/proxy_gui.py

@@ -37,7 +37,7 @@ from hippolyzer.lib.base.message.msgtypes import MsgType
 from hippolyzer.lib.base.message.template_dict import TemplateDictionary
 from hippolyzer.lib.base.ui_helpers import loadUi
 import hippolyzer.lib.base.serialization as se
-from hippolyzer.lib.base.network.transport import Direction, WrappingUDPTransport
+from hippolyzer.lib.base.network.transport import Direction, SocketUDPTransport
 from hippolyzer.lib.proxy.addons import BaseInteractionManager, AddonManager
 from hippolyzer.lib.proxy.ca_utils import setup_ca_everywhere
 from hippolyzer.lib.proxy.caps_client import ProxyCapsClient
@@ -646,7 +646,7 @@ class MessageBuilderWindow(QtWidgets.QMainWindow):
             transport = None
             off_circuit = self.checkOffCircuit.isChecked()
             if off_circuit:
-                transport = WrappingUDPTransport(socket.socket(socket.AF_INET, socket.SOCK_DGRAM))
+                transport = SocketUDPTransport(socket.socket(socket.AF_INET, socket.SOCK_DGRAM))
             region.circuit.send_message(msg, transport=transport)
             if off_circuit:
                 transport.close()

hippolyzer/lib/base/network/transport.py

@@ -58,7 +58,7 @@ class AbstractUDPTransport(abc.ABC):
         pass
 
 
-class WrappingUDPTransport(AbstractUDPTransport):
+class SocketUDPTransport(AbstractUDPTransport):
     def __init__(self, transport: Union[asyncio.DatagramTransport, socket.socket]):
         super().__init__()
         self.transport = transport

hippolyzer/lib/proxy/addon_utils.py

@@ -181,13 +181,15 @@ class BaseAddon(abc.ABC):
     def handle_region_changed(self, session: Session, region: ProxiedRegion):
         pass
 
+    def handle_circuit_created(self, session: Session, region: ProxiedRegion):
+        pass
+
     def handle_rlv_command(self, session: Session, region: ProxiedRegion, source: UUID,
                            cmd: str, options: List[str], param: str):
         pass
 
     def handle_proxied_packet(self, session_manager: SessionManager, packet: UDPPacket,
-                              session: Optional[Session], region: Optional[ProxiedRegion],
-                              message: Optional[Message]):
+                              session: Optional[Session], region: Optional[ProxiedRegion]):
         pass
 
 

hippolyzer/lib/proxy/addons.py

@@ -16,15 +16,15 @@ from types import ModuleType
 from typing import *
 
 from hippolyzer.lib.base.datatypes import UUID
+from hippolyzer.lib.base.message.message import Message
+from hippolyzer.lib.base.network.transport import UDPPacket
 from hippolyzer.lib.proxy import addon_ctx
 from hippolyzer.lib.proxy.task_scheduler import TaskLifeScope, TaskScheduler
 
 if TYPE_CHECKING:
     from hippolyzer.lib.proxy.commands import CommandDetails, WrappedCommandCallable
     from hippolyzer.lib.proxy.http_flow import HippoHTTPFlow
-    from hippolyzer.lib.base.message.message import Message
-    from hippolyzer.lib.proxy.objects import Object
-    from hippolyzer.lib.base.network.transport import UDPPacket
+    from hippolyzer.lib.proxy.object_manager import Object
     from hippolyzer.lib.proxy.region import ProxiedRegion
     from hippolyzer.lib.proxy.sessions import Session, SessionManager
 
@@ -526,9 +526,14 @@ class AddonManager:
         with addon_ctx.push(session, region):
             return cls._call_all_addon_hooks("handle_region_changed", session, region)
 
+    @classmethod
+    def handle_circuit_created(cls, session: Session, region: ProxiedRegion):
+        with addon_ctx.push(session, region):
+            return cls._call_all_addon_hooks("handle_circuit_created", session, region)
+
     @classmethod
     def handle_proxied_packet(cls, session_manager: SessionManager, packet: UDPPacket,
-                              session: Optional[Session], region: Optional[ProxiedRegion],
-                              message: Optional[Message]):
-        return cls._call_all_addon_hooks("handle_proxied_packet", session_manager,
-                                         packet, session, region, message)
+                              session: Optional[Session], region: Optional[ProxiedRegion]):
+        with addon_ctx.push(session, region):
+            return cls._call_all_addon_hooks("handle_proxied_packet", session_manager,
+                                             packet, session, region)

hippolyzer/lib/proxy/lludp_proxy.py

@@ -44,29 +44,20 @@ class InterceptingLLUDPProxyProtocol(UDPProxyProtocol):
             raise PermissionError(f"UDPBanned message {msg.name}")
 
     def _handle_proxied_packet(self, packet: UDPPacket):
-        message: Optional[Message] = None
         region: Optional[ProxiedRegion] = None
         # Try to do an initial region lookup so we have it for handle_proxied_packet()
         if self.session:
             region = self.session.region_by_circuit_addr(packet.far_addr)
-        deserialize_exc = None
-        try:
-            message = self.deserializer.deserialize(packet.data)
-            message.direction = packet.direction
-            message.sender = packet.src_addr
-        except Exception as e:
-            # Hang onto this since handle_proxied_packet doesn't need a parseable
-            # message. If that hook doesn't handle the packet then re-raise.
-            deserialize_exc = e
 
+        # the proxied packet handler is allowed to mutate `packet.data` before
+        # the message gets parsed.
         if AddonManager.handle_proxied_packet(self.session_manager, packet,
-                                              self.session, region, message):
-            # Swallow any error raised by above message deserialization, it was handled.
+                                              self.session, region):
             return
 
-        if deserialize_exc is not None:
-            # handle_proxied_packet() didn't deal with the error, so it's fatal.
-            raise deserialize_exc
+        message = self.deserializer.deserialize(packet.data)
+        message.direction = packet.direction
+        message.sender = packet.src_addr
 
         assert message is not None
         # Check for UDP bans on inbound messages

hippolyzer/lib/proxy/sessions.py

@@ -13,6 +13,7 @@ from hippolyzer.lib.base.datatypes import UUID
 from hippolyzer.lib.base.message.message import Message
 from hippolyzer.lib.base.message.message_handler import MessageHandler
 from hippolyzer.lib.client.state import BaseClientSession
+from hippolyzer.lib.proxy.addons import AddonManager
 from hippolyzer.lib.proxy.circuit import ProxiedCircuit
 from hippolyzer.lib.proxy.http_asset_repo import HTTPAssetRepo
 from hippolyzer.lib.proxy.http_proxy import HTTPFlowContext, is_asset_server_cap_name, SerializedCapData
@@ -136,6 +137,7 @@ class Session(BaseClientSession):
                         )
                     region.circuit = ProxiedCircuit(
                         near_addr, circuit_addr, transport, logging_hook=logging_hook)
+                    AddonManager.handle_circuit_created(self, region)
                     return True
                 if region.circuit and region.circuit.is_alive:
                     # Whatever, already open

hippolyzer/lib/proxy/transport.py

@@ -1,10 +1,10 @@
 import socket
 import struct
 
-from hippolyzer.lib.base.network.transport import WrappingUDPTransport, UDPPacket
+from hippolyzer.lib.base.network.transport import SocketUDPTransport, UDPPacket
 
 
-class SOCKS5UDPTransport(WrappingUDPTransport):
+class SOCKS5UDPTransport(SocketUDPTransport):
     HEADER_STRUCT = struct.Struct("!HBB4sH")
 
     @classmethod