diff options
Diffstat (limited to 'libmproxy/protocol/primitives.py')
-rw-r--r-- | libmproxy/protocol/primitives.py | 74 |
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 |