split up protocol\__init__.py

1 files changed, 73 insertions, 1 deletions

diff --git a/libmproxy/protocol/primitives.py b/libmproxy/protocol/primitives.py
index f2701458..5b95e9e5 100644
--- a/libmproxy/protocol/primitives.py
+++ b/libmproxy/protocol/primitives.py
@@ -1,8 +1,12 @@
 from .. import stateobject, utils, version
+from ..proxy.primitives import AddressPriority
 from ..proxy.connection import ClientConnection, ServerConnection
 import copy
 
 
+KILL = 0  # const for killed requests
+
+
 class BackreferenceMixin(object):
     """
     If an attribute from the _backrefattr tuple is set,
@@ -127,4 +131,72 @@ class Flow(stateobject.SimpleStateObject, BackreferenceMixin):
         """
         if self._backup:
             self._load_state(self._backup)
-            self._backup = None
\ No newline at end of file
+            self._backup = None
+
+
+class ProtocolHandler(object):
+    def __init__(self, c):
+        self.c = c
+        """@type: libmproxy.proxy.ConnectionHandler"""
+
+    def handle_messages(self):
+        """
+        This method gets called if a client connection has been made. Depending on the proxy settings,
+        a server connection might already exist as well.
+        """
+        raise NotImplementedError  # pragma: nocover
+
+    def handle_error(self, error):
+        """
+        This method gets called should there be an uncaught exception during the connection.
+        This might happen outside of handle_messages, e.g. if the initial SSL handshake fails in transparent mode.
+        """
+        raise error  # pragma: nocover
+
+
+class TemporaryServerChangeMixin(object):
+    """
+    This mixin allows safe modification of the target server,
+    without any need to expose the ConnectionHandler to the Flow.
+    """
+    def change_server(self, address, ssl):
+        if address == self.c.server_conn.address():
+            return
+        priority = AddressPriority.MANUALLY_CHANGED
+
+        if self.c.server_conn.priority > priority:
+            self.log("Attempt to change server address, "
+                     "but priority is too low (is: %s, got: %s)" % (self.server_conn.priority, priority))
+            return
+
+        self.log("Temporarily change server connection: %s:%s -> %s:%s" % (
+            self.c.server_conn.address.host,
+            self.c.server_conn.address.port,
+            address.host,
+            address.port
+        ))
+
+        if not hasattr(self, "_backup_server_conn"):
+            self._backup_server_conn = self.c.server_conn
+            self.c.server_conn = None
+        else:  # This is at least the second temporary change. We can kill the current connection.
+            self.c.del_server_connection()
+
+        self.c.set_server_address(address, priority)
+        if ssl:
+            self.c.establish_ssl(server=True)
+
+    def restore_server(self):
+        if not hasattr(self, "_backup_server_conn"):
+            return
+
+        self.log("Restore original server connection: %s:%s -> %s:%s" % (
+            self.c.server_conn.address.host,
+            self.c.server_conn.address.port,
+            self._backup_server_conn.host,
+            self._backup_server_conn.port
+        ))
+
+        self.c.del_server_connection()
+        self.c.server_conn = self._backup_server_conn
+        del self._backup_server_conn
\ No newline at end of file