diff --git a/examples/change_upstream_proxy.py b/examples/change_upstream_proxy.py index 7782dd848..8f58e1f21 100644 --- a/examples/change_upstream_proxy.py +++ b/examples/change_upstream_proxy.py @@ -1,29 +1,34 @@ # This scripts demonstrates how mitmproxy can switch to a second/different upstream proxy # in upstream proxy mode. # -# Usage: mitmdump -U http://default-upstream-proxy.local:8080/ -s -# "change_upstream_proxy.py host" -from libmproxy.protocol.http import send_connect_request - -alternative_upstream_proxy = ("localhost", 8082) +# 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 should_redirect(flow): - return flow.request.host == "example.com" +def proxy_address(flow): + # Poor man's loadbalancing: route every second domain through the alternative proxy. + if hash(flow.request.host) % 2 == 1: + return ("localhost", 8082) + else: + return ("localhost", 8081) def request(context, flow): - if flow.live and should_redirect(flow): - - # If you want to change the target server, you should modify flow.request.host and flow.request.port - # flow.live.change_server should only be used by inline scripts to change the upstream proxy, - # unless you are sure that you know what you are doing. - server_changed = flow.live.change_server( - alternative_upstream_proxy, - persistent_change=True) - if flow.request.scheme == "https" and server_changed: - send_connect_request( - flow.live.c.server_conn, - flow.request.host, - flow.request.port) - flow.live.c.establish_ssl(server=True) + if flow.request.method == "CONNECT": + # If the decision is done by domain, one could also modify the server address here. + # We do it after CONNECT here to have the request data available as well. + 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) diff --git a/libmproxy/protocol/base.py b/libmproxy/protocol/base.py index 1c9b356c0..3440cb018 100644 --- a/libmproxy/protocol/base.py +++ b/libmproxy/protocol/base.py @@ -116,6 +116,10 @@ class ServerConnectionMixin(object): self._disconnect() self.log("Set new server address: " + repr(address), "debug") self.server_conn.address = address + 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) diff --git a/libmproxy/protocol/http.py b/libmproxy/protocol/http.py index 345b3aa84..3b62c3897 100644 --- a/libmproxy/protocol/http.py +++ b/libmproxy/protocol/http.py @@ -207,6 +207,9 @@ class ConnectServerConnection(object): def __getattr__(self, item): return getattr(self.via, item) + def __nonzero__(self): + return bool(self.via) + class UpstreamConnectLayer(Layer): def __init__(self, ctx, connect_request): @@ -221,19 +224,22 @@ class UpstreamConnectLayer(Layer): layer = self.ctx.next_layer(self) layer() + def _send_connect_request(self): + self.send_request(self.connect_request) + resp = self.read_response("CONNECT") + if resp.code != 200: + raise ProtocolException("Reconnect: Upstream server refuses CONNECT request") + def connect(self): if not self.server_conn: self.ctx.connect() - self.send_request(self.connect_request) + self._send_connect_request() else: pass # swallow the message def reconnect(self): self.ctx.reconnect() - self.send_request(self.connect_request) - resp = self.read_response("CONNECT") - if resp.code != 200: - raise ProtocolException("Reconnect: Upstream server refuses CONNECT request") + self._send_connect_request() def set_server(self, address, server_tls=None, sni=None, depth=1): if depth == 1: @@ -386,7 +392,7 @@ class HttpLayer(Layer): if self.supports_streaming: flow.response = self.read_response_headers() else: - flow.response = self.read_response() + flow.response = self.read_response(flow.request.method) try: get_response() @@ -473,13 +479,6 @@ class HttpLayer(Layer): # Establish connection is neccessary. if not self.server_conn: self.connect() - - # SetServer is not guaranteed to work with TLS: - # If there's not TlsLayer below which could catch the exception, - # TLS will not be established. - if tls and not self.server_conn.tls_established: - raise ProtocolException( - "Cannot upgrade to SSL, no TLS layer on the protocol stack.") else: if not self.server_conn: self.connect() diff --git a/libmproxy/protocol/tls.py b/libmproxy/protocol/tls.py index b85a6595e..2646ec4f2 100644 --- a/libmproxy/protocol/tls.py +++ b/libmproxy/protocol/tls.py @@ -152,10 +152,12 @@ class TlsLayer(Layer): self._establish_tls_with_server() def set_server(self, address, server_tls=None, sni=None, depth=1): - self.ctx.set_server(address, server_tls, sni, depth) if depth == 1 and server_tls is not None: + self.ctx.set_server(address, None, None, 1) self._sni_from_server_change = sni self._server_tls = server_tls + else: + self.ctx.set_server(address, server_tls, sni, depth) @property def sni_for_server_connection(self):