From af17525071d4ef193afc3f1369098229bbfb8149 Mon Sep 17 00:00:00 2001
From: Salad Dais <SaladDais@users.noreply.github.com>
Date: Wed, 2 Jun 2021 11:25:53 +0000
Subject: [PATCH] Remove Circuit dependency on parent Region

---
 hippolyzer/lib/proxy/circuit.py  | 17 ++++++-----------
 hippolyzer/lib/proxy/sessions.py | 10 +++++++++-
 2 files changed, 15 insertions(+), 12 deletions(-)

diff --git a/hippolyzer/lib/proxy/circuit.py b/hippolyzer/lib/proxy/circuit.py
index 8758c4c..d4faa29 100644
--- a/hippolyzer/lib/proxy/circuit.py
+++ b/hippolyzer/lib/proxy/circuit.py
@@ -12,13 +12,12 @@ from hippolyzer.lib.base.message.udpserializer import UDPMessageSerializer
 from hippolyzer.lib.proxy.packets import Direction, ProxiedUDPPacket
 from hippolyzer.lib.proxy.message import ProxiedMessage
 
-if TYPE_CHECKING:
-    from hippolyzer.lib.proxy.region import ProxiedRegion
-    from hippolyzer.lib.proxy.message_logger import BaseMessageLogger
+
+LLUDP_LOGGING_HOOK = Optional[Callable[[ProxiedMessage], Any]]
 
 
 class ProxiedCircuit:
-    def __init__(self, near_host, far_host, transport, region: Optional[ProxiedRegion] = None,
+    def __init__(self, near_host, far_host, transport, logging_hook: LLUDP_LOGGING_HOOK = None,
                  socks_transport: Optional[bool] = None):
         self.near_host = near_host
         self.host = far_host
@@ -29,11 +28,7 @@ class ProxiedCircuit:
         self.out_injections = InjectionTracker(0)
         self.serializer = UDPMessageSerializer()
         self.last_packet_at = dt.datetime.now()
-        self.region: Optional[ProxiedRegion] = region
-        message_logger = None
-        if region:
-            message_logger = region.session().session_manager.message_logger
-        self.message_logger: Optional[BaseMessageLogger] = message_logger
+        self.logging_hook: LLUDP_LOGGING_HOOK = logging_hook
 
     def _send_prepared_message(self, message: ProxiedMessage, direction, transport=None):
         try:
@@ -41,8 +36,8 @@ class ProxiedCircuit:
         except:
             logging.exception(f"Failed to serialize: {message.to_dict()!r}")
             raise
-        if self.message_logger and message.injected:
-            self.message_logger.log_lludp_message(self.region.session(), self.region, message)
+        if self.logging_hook and message.injected:
+            self.logging_hook(message)
         return self.send_datagram(serialized, direction, transport=transport)
 
     def send_datagram(self, data: bytes, direction: Direction, transport=None):
diff --git a/hippolyzer/lib/proxy/sessions.py b/hippolyzer/lib/proxy/sessions.py
index 8b1054c..db91a1f 100644
--- a/hippolyzer/lib/proxy/sessions.py
+++ b/hippolyzer/lib/proxy/sessions.py
@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import dataclasses
 import datetime
+import functools
 import logging
 import multiprocessing
 import weakref
@@ -124,8 +125,15 @@ class Session:
         for region in self.regions:
             if region.circuit_addr == circuit_addr:
                 if not region.circuit or not region.circuit.is_alive:
+                    logging_hook = None
+                    if self.session_manager.message_logger:
+                        logging_hook = functools.partial(
+                            self.session_manager.message_logger.log_lludp_message,
+                            self,
+                            region,
+                        )
                     region.circuit = ProxiedCircuit(
-                        near_addr, circuit_addr, transport, region=region)
+                        near_addr, circuit_addr, transport, logging_hook=logging_hook)
                     return True
                 if region.circuit and region.circuit.is_alive:
                     # Whatever, already open