Use an enum for http protocol modes

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=()):
         """