From 5be35d258fd95f9f99ee6e8a32dc234f6b868385 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Sat, 12 Nov 2016 10:06:57 +1300 Subject: Use an enum for http protocol modes --- mitmproxy/proxy/protocol/http.py | 33 +++++++++++++++++---------------- mitmproxy/proxy/root_context.py | 11 ++++++----- 2 files changed, 23 insertions(+), 21 deletions(-) diff --git a/mitmproxy/proxy/protocol/http.py b/mitmproxy/proxy/protocol/http.py index 3bc33ab0..e974ffe2 100644 --- a/mitmproxy/proxy/protocol/http.py +++ b/mitmproxy/proxy/protocol/http.py @@ -1,6 +1,8 @@ import h2.exceptions import time import traceback +import enum + from mitmproxy import exceptions from mitmproxy import http from mitmproxy import flow @@ -112,13 +114,18 @@ class UpstreamConnectLayer(base.Layer): self.server_conn.address = address +class HTTPMode(enum.Enum): + regular = 1 + transparent = 2 + upstream = 3 + + FIRSTLINES = set(["absolute", "relative", "authority"]) # At this point, we see only a subset of the proxy modes -MODES = set(["regular", "transparent", "upstream"]) MODE_REQUEST_FORMS = { - "regular": ("authority", "absolute"), - "transparent": ("relative"), - "upstream": ("authority", "absolute"), + HTTPMode.regular: ("authority", "absolute"), + HTTPMode.transparent: ("relative"), + HTTPMode.upstream: ("authority", "absolute"), } @@ -137,8 +144,6 @@ class HttpLayer(base.Layer): def __init__(self, ctx, mode): super().__init__(ctx) - if mode not in MODES: - raise exceptions.ServerException("Invalid mode: %s" % mode) self.mode = mode self.__initial_server_conn = None "Contains the original destination in transparent mode, which needs to be restored" @@ -150,7 +155,7 @@ class HttpLayer(base.Layer): self.connect_request = False def __call__(self): - if self.mode == "transparent": + if self.mode == HTTPMode.transparent: self.__initial_server_tls = self.server_tls self.__initial_server_conn = self.server_conn while True: @@ -199,7 +204,7 @@ class HttpLayer(base.Layer): try: # Regular Proxy Mode: Handle CONNECT - if self.mode == "regular" and request.first_line_format == "authority": + if self.mode is HTTPMode.regular and request.first_line_format == "authority": self.connect_request = True self.set_server((request.host, request.port)) self.send_response(http.make_connect_response(request.data.http_version)) @@ -219,7 +224,7 @@ class HttpLayer(base.Layer): f.request.headers["Host"] = self.config.upstream_server.address.host # set upstream auth - if self.mode == "upstream" and self.config.upstream_auth is not None: + if self.mode is HTTPMode.upstream and self.config.upstream_auth is not None: f.request.headers["Proxy-Authorization"] = self.config.upstream_auth # Determine .scheme, .host and .port attributes for inline scripts. @@ -227,11 +232,7 @@ class HttpLayer(base.Layer): # For authority-form requests, we only need to determine the request scheme. # For relative-form requests, we need to determine host and port as # well. - if self.mode == "regular": - pass # only absolute-form at this point, nothing to do here. - elif self.mode == "upstream": - pass - else: + if self.mode is HTTPMode.transparent: # Setting request.host also updates the host header, which we want to preserve host_header = f.request.headers.get("host", None) f.request.host = self.__initial_server_conn.address.host @@ -395,7 +396,7 @@ class HttpLayer(base.Layer): address = tcp.Address((host, port)) tls = (scheme == "https") - if self.mode == "regular" or self.mode == "transparent": + if self.mode is HTTPMode.regular or self.mode is HTTPMode.transparent: # If there's an existing connection that doesn't match our expectations, kill it. if address != self.server_conn.address or tls != self.server_tls: self.set_server(address) @@ -414,7 +415,7 @@ class HttpLayer(base.Layer): if self.config.authenticator.authenticate(request.headers): self.config.authenticator.clean(request.headers) else: - if self.mode == "transparent": + if self.mode == HTTPMode.transparent: self.send_response(http.make_error_response( 401, "Authentication Required", diff --git a/mitmproxy/proxy/root_context.py b/mitmproxy/proxy/root_context.py index 4362347b..50dbe79e 100644 --- a/mitmproxy/proxy/root_context.py +++ b/mitmproxy/proxy/root_context.py @@ -2,6 +2,7 @@ from mitmproxy import log from mitmproxy import exceptions from mitmproxy.proxy import protocol from mitmproxy.proxy import modes +from mitmproxy.proxy.protocol import http class RootContext: @@ -70,9 +71,9 @@ class RootContext: # 3. In Http Proxy mode and Upstream Proxy mode, the next layer is fixed. if isinstance(top_layer, protocol.TlsLayer): if isinstance(top_layer.ctx, modes.HttpProxy): - return protocol.Http1Layer(top_layer, "regular") + return protocol.Http1Layer(top_layer, http.HTTPMode.regular) if isinstance(top_layer.ctx, modes.HttpUpstreamProxy): - return protocol.Http1Layer(top_layer, "upstream") + return protocol.Http1Layer(top_layer, http.HTTPMode.upstream) # 4. Check for other TLS cases (e.g. after CONNECT). if client_tls: @@ -86,12 +87,12 @@ class RootContext: if isinstance(top_layer, protocol.TlsLayer): alpn = top_layer.client_conn.get_alpn_proto_negotiated() if alpn == b'h2': - return protocol.Http2Layer(top_layer, 'transparent') + return protocol.Http2Layer(top_layer, http.HTTPMode.transparent) if alpn == b'http/1.1': - return protocol.Http1Layer(top_layer, 'transparent') + return protocol.Http1Layer(top_layer, http.HTTPMode.transparent) # 6. Assume HTTP1 by default - return protocol.Http1Layer(top_layer, 'transparent') + return protocol.Http1Layer(top_layer, http.HTTPMode.transparent) def log(self, msg, level, subs=()): """ -- cgit v1.2.3