Merge pull request #3526 from pierlon/feature/allow-hosts

Add --allow_hosts option

5 files changed, 33 insertions, 15 deletions

diff --git a/mitmproxy/options.py b/mitmproxy/options.py
index a6ab3d50..56146153 100644
--- a/mitmproxy/options.py
+++ b/mitmproxy/options.py
@@ -68,6 +68,10 @@ class Options(optmanager.OptManager):
             """
         )
         self.add_option(
+            "allow_hosts", Sequence[str], [],
+            "Opposite of --ignore-hosts."
+        )
+        self.add_option(
             "listen_host", str, "",
             "Address to bind proxy to."
         )
diff --git a/mitmproxy/proxy/config.py b/mitmproxy/proxy/config.py
index f32d3086..75e372ae 100644
--- a/mitmproxy/proxy/config.py
+++ b/mitmproxy/proxy/config.py
@@ -14,7 +14,8 @@ CONF_BASENAME = "mitmproxy"
 
 class HostMatcher:
 
-    def __init__(self, patterns=tuple()):
+    def __init__(self, handle, patterns=tuple()):
+        self.handle = handle
         self.patterns = list(patterns)
         self.regexes = [re.compile(p, re.IGNORECASE) for p in self.patterns]
 
@@ -22,10 +23,10 @@ class HostMatcher:
         if not address:
             return False
         host = "%s:%s" % address
-        if any(rex.search(host) for rex in self.regexes):
-            return True
-        else:
-            return False
+        if self.handle in ["ignore", "tcp"]:
+            return any(rex.search(host) for rex in self.regexes)
+        else:  # self.handle == "allow"
+            return any(not rex.search(host) for rex in self.regexes)
 
     def __bool__(self):
         return bool(self.patterns)
@@ -36,7 +37,7 @@ class ProxyConfig:
     def __init__(self, options: moptions.Options) -> None:
         self.options = options
 
-        self.check_ignore: HostMatcher = None
+        self.check_filter: HostMatcher = None
         self.check_tcp: HostMatcher = None
         self.certstore: certs.CertStore = None
         self.upstream_server: typing.Optional[server_spec.ServerSpec] = None
@@ -44,10 +45,18 @@ class ProxyConfig:
         options.changed.connect(self.configure)
 
     def configure(self, options: moptions.Options, updated: typing.Any) -> None:
-        if "ignore_hosts" in updated:
-            self.check_ignore = HostMatcher(options.ignore_hosts)
+        if options.allow_hosts and options.ignore_hosts:
+            raise exceptions.OptionsError("--ignore-hosts and --allow-hosts are mutually "
+                                          "exclusive; please choose one.")
+
+        if options.ignore_hosts:
+            self.check_filter = HostMatcher("ignore", options.ignore_hosts)
+        elif options.allow_hosts:
+            self.check_filter = HostMatcher("allow", options.allow_hosts)
+        else:
+            self.check_filter = HostMatcher(False)
         if "tcp_hosts" in updated:
-            self.check_tcp = HostMatcher(options.tcp_hosts)
+            self.check_tcp = HostMatcher("tcp", options.tcp_hosts)
 
         certstore_path = os.path.expanduser(options.confdir)
         if not os.path.exists(os.path.dirname(certstore_path)):
diff --git a/mitmproxy/proxy/root_context.py b/mitmproxy/proxy/root_context.py
index eb0008cf..4805f874 100644
--- a/mitmproxy/proxy/root_context.py
+++ b/mitmproxy/proxy/root_context.py
@@ -48,17 +48,17 @@ class RootContext:
             raise exceptions.ProtocolException(str(e))
         client_tls = tls.is_tls_record_magic(d)
 
-        # 1. check for --ignore
-        if self.config.check_ignore:
-            ignore = self.config.check_ignore(top_layer.server_conn.address)
-            if not ignore and client_tls:
+        # 1. check for filter
+        if self.config.check_filter:
+            is_filtered = self.config.check_filter(top_layer.server_conn.address)
+            if not is_filtered and client_tls:
                 try:
                     client_hello = tls.ClientHello.from_file(self.client_conn.rfile)
                 except exceptions.TlsProtocolException as e:
                     self.log("Cannot parse Client Hello: %s" % repr(e), "error")
                 else:
-                    ignore = self.config.check_ignore((client_hello.sni, 443))
-            if ignore:
+                    is_filtered = self.config.check_filter((client_hello.sni, 443))
+            if is_filtered:
                 return protocol.RawTCPLayer(top_layer, ignore=True)
 
         # 2. Always insert a TLS layer, even if there's neither client nor server tls.
diff --git a/mitmproxy/tools/cmdline.py b/mitmproxy/tools/cmdline.py
index eb4a984d..2c7817fa 100644
--- a/mitmproxy/tools/cmdline.py
+++ b/mitmproxy/tools/cmdline.py
@@ -57,6 +57,7 @@ def common_options(parser, opts):
     opts.make_parser(group, "listen_port", metavar="PORT", short="p")
     opts.make_parser(group, "server", short="n")
     opts.make_parser(group, "ignore_hosts", metavar="HOST")
+    opts.make_parser(group, "allow_hosts", metavar="HOST")
     opts.make_parser(group, "tcp_hosts", metavar="HOST")
     opts.make_parser(group, "upstream_auth", metavar="USER:PASS")
     opts.make_parser(group, "proxyauth", metavar="SPEC")
diff --git a/mitmproxy/tools/console/statusbar.py b/mitmproxy/tools/console/statusbar.py
index 2d32f487..56f0674f 100644
--- a/mitmproxy/tools/console/statusbar.py
+++ b/mitmproxy/tools/console/statusbar.py
@@ -215,6 +215,10 @@ class StatusBar(urwid.WidgetWrap):
             r.append("[")
             r.append(("heading_key", "I"))
             r.append("gnore:%d]" % len(self.master.options.ignore_hosts))
+        elif self.master.options.allow_hosts:
+            r.append("[")
+            r.append(("heading_key", "A"))
+            r.append("llow:%d]" % len(self.master.options.allow_hosts))
         if self.master.options.tcp_hosts:
             r.append("[")
             r.append(("heading_key", "T"))