diff options
| -rw-r--r-- | examples/change_upstream_proxy.py | 12 | ||||
| -rw-r--r-- | libmproxy/protocol/base.py | 23 | ||||
| -rw-r--r-- | libmproxy/protocol/http.py | 32 | ||||
| -rw-r--r-- | libmproxy/protocol/tls.py | 8 |
4 files changed, 36 insertions, 39 deletions
diff --git a/examples/change_upstream_proxy.py b/examples/change_upstream_proxy.py index 8f58e1f2..9c454897 100644 --- a/examples/change_upstream_proxy.py +++ b/examples/change_upstream_proxy.py @@ -4,7 +4,6 @@ # Usage: mitmdump -U http://default-upstream-proxy.local:8080/ -s change_upstream_proxy.py # # If you want to change the target server, you should modify flow.request.host and flow.request.port -# flow.live.set_server should only be used by inline scripts to change the upstream proxy. def proxy_address(flow): @@ -22,13 +21,4 @@ def request(context, flow): return address = proxy_address(flow) if flow.live: - if flow.request.scheme == "http": - # For a normal HTTP request, we just change the proxy server and we're done! - if address != flow.live.server_conn.address: - flow.live.set_server(address, depth=1) - else: - # If we have CONNECTed (and thereby established "destination state"), the story is - # a bit more complex. Now we don't want to change the top level address (which is - # the connect destination) but the address below that. (Notice the `.via` and depth=2). - if address != flow.live.server_conn.via.address: - flow.live.set_server(address, depth=2) + flow.live.change_upstream_proxy_server(address)
\ No newline at end of file diff --git a/libmproxy/protocol/base.py b/libmproxy/protocol/base.py index f1718065..f27cb04b 100644 --- a/libmproxy/protocol/base.py +++ b/libmproxy/protocol/base.py @@ -116,19 +116,16 @@ class ServerConnectionMixin(object): "The proxy shall not connect to itself.".format(repr(address)) ) - def set_server(self, address, server_tls=None, sni=None, depth=1): - if depth == 1: - if self.server_conn: - self.disconnect() - self.log("Set new server address: " + repr(address), "debug") - self.server_conn.address = address - self.__check_self_connect() - if server_tls: - raise ProtocolException( - "Cannot upgrade to TLS, no TLS layer on the protocol stack." - ) - else: - self.ctx.set_server(address, server_tls, sni, depth - 1) + def set_server(self, address, server_tls=None, sni=None): + if self.server_conn: + self.disconnect() + self.log("Set new server address: " + repr(address), "debug") + self.server_conn.address = address + self.__check_self_connect() + if server_tls: + raise ProtocolException( + "Cannot upgrade to TLS, no TLS layer on the protocol stack." + ) def disconnect(self): """ diff --git a/libmproxy/protocol/http.py b/libmproxy/protocol/http.py index 3c934393..f2265c34 100644 --- a/libmproxy/protocol/http.py +++ b/libmproxy/protocol/http.py @@ -304,16 +304,22 @@ class UpstreamConnectLayer(Layer): else: pass # swallow the message - def set_server(self, address, server_tls=None, sni=None, depth=1): - if depth == 1: - if self.ctx.server_conn: - self.ctx.disconnect() - address = Address.wrap(address) - self.connect_request.host = address.host - self.connect_request.port = address.port - self.server_conn.address = address - else: - self.ctx.set_server(address, server_tls, sni, depth - 1) + def change_upstream_proxy_server(self, address): + if address != self.server_conn.via.address: + self.ctx.set_server(address) + + def set_server(self, address, server_tls=None, sni=None): + if self.ctx.server_conn: + self.ctx.disconnect() + address = Address.wrap(address) + self.connect_request.host = address.host + self.connect_request.port = address.port + self.server_conn.address = address + + if server_tls: + raise ProtocolException( + "Cannot upgrade to TLS, no TLS layer on the protocol stack." + ) class HttpLayer(Layer): @@ -388,6 +394,12 @@ class HttpLayer(Layer): finally: flow.live = False + def change_upstream_proxy_server(self, address): + # Make set_upstream_proxy_server always available, + # even if there's no UpstreamConnectLayer + if address != self.server_conn.address: + return self.set_server(address) + def handle_regular_mode_connect(self, request): self.set_server((request.host, request.port)) self.send_response(make_connect_response(request.httpversion)) diff --git a/libmproxy/protocol/tls.py b/libmproxy/protocol/tls.py index 00e016ea..a62b1a22 100644 --- a/libmproxy/protocol/tls.py +++ b/libmproxy/protocol/tls.py @@ -338,13 +338,11 @@ class TlsLayer(Layer): if self._server_tls and not self.server_conn.tls_established: self._establish_tls_with_server() - def set_server(self, address, server_tls=None, sni=None, depth=1): - if depth == 1 and server_tls is not None: - self.ctx.set_server(address, None, None, 1) + def set_server(self, address, server_tls=None, sni=None): + if server_tls is not None: self._sni_from_server_change = sni self._server_tls = server_tls - else: - self.ctx.set_server(address, server_tls, sni, depth) + self.ctx.set_server(address, None, None) @property def sni_for_server_connection(self): |