add generic tcp proxying, fix #374

2 files changed, 30 insertions, 17 deletions

diff --git a/libmproxy/proxy/config.py b/libmproxy/proxy/config.py
index abdb7c41..948decc1 100644
--- a/libmproxy/proxy/config.py
+++ b/libmproxy/proxy/config.py
@@ -1,7 +1,7 @@
 from __future__ import absolute_import
 import os
 import re
-from netlib import http_auth, certutils
+from netlib import http_auth, certutils, tcp
 from .. import utils, platform, version
 from .primitives import RegularProxyMode, TransparentProxyMode, UpstreamProxyMode, ReverseProxyMode, Socks5ProxyMode
 
@@ -10,8 +10,21 @@ CONF_BASENAME = "mitmproxy"
 CONF_DIR = "~/.mitmproxy"
 
 
-def parse_host_pattern(patterns):
-    return [re.compile(p, re.IGNORECASE) for p in patterns]
+class HostMatcher(object):
+    def __init__(self, patterns=[]):
+        self.patterns = list(patterns)
+        self.regexes = [re.compile(p, re.IGNORECASE) for p in self.patterns]
+
+    def __call__(self, address):
+        address = tcp.Address.wrap(address)
+        host = "%s:%s" % (address.host, address.port)
+        if any(rex.search(host) for rex in self.regexes):
+            return True
+        else:
+            return False
+
+    def __nonzero__(self):
+        return bool(self.patterns)
 
 
 class ProxyConfig:
@@ -19,7 +32,7 @@ class ProxyConfig:
                  confdir=CONF_DIR, clientcerts=None,
                  no_upstream_cert=False, body_size_limit=None,
                  mode=None, upstream_server=None, http_form_in=None, http_form_out=None,
-                 authenticator=None, ignore=[],
+                 authenticator=None, ignore_hosts=[], tcp_hosts=[],
                  ciphers=None, certs=[], certforward=False, ssl_ports=TRANSPARENT_SSL_PORTS):
         self.host = host
         self.port = port
@@ -44,7 +57,8 @@ class ProxyConfig:
         self.mode.http_form_in = http_form_in or self.mode.http_form_in
         self.mode.http_form_out = http_form_out or self.mode.http_form_out
 
-        self.ignore = parse_host_pattern(ignore)
+        self.check_ignore = HostMatcher(ignore_hosts)
+        self.check_tcp = HostMatcher(tcp_hosts)
         self.authenticator = authenticator
         self.confdir = os.path.expanduser(confdir)
         self.certstore = certutils.CertStore.from_store(self.confdir, CONF_BASENAME)
@@ -124,7 +138,8 @@ def process_proxy_options(parser, options):
         upstream_server=upstream_server,
         http_form_in=options.http_form_in,
         http_form_out=options.http_form_out,
-        ignore=options.ignore,
+        ignore_hosts=options.ignore_hosts,
+        tcp_hosts=options.tcp_hosts,
         authenticator=authenticator,
         ciphers=options.ciphers,
         certs=certs,
diff --git a/libmproxy/proxy/server.py b/libmproxy/proxy/server.py
index 4c7fbbf0..fdf6405a 100644
--- a/libmproxy/proxy/server.py
+++ b/libmproxy/proxy/server.py
@@ -70,13 +70,15 @@ class ConnectionHandler:
 
             # Can we already identify the target server and connect to it?
             client_ssl, server_ssl = False, False
+            conn_kwargs = dict()
             upstream_info = self.config.mode.get_upstream_server(self.client_conn)
             if upstream_info:
                 self.set_server_address(upstream_info[2:])
                 client_ssl, server_ssl = upstream_info[:2]
-                if self.check_ignore_address(self.server_conn.address):
+                if self.config.check_ignore(self.server_conn.address):
                     self.log("Ignore host: %s:%s" % self.server_conn.address(), "info")
                     self.conntype = "tcp"
+                    conn_kwargs["log"] = False
                     client_ssl, server_ssl = False, False
             else:
                 pass  # No upstream info from the metadata: upstream info in the protocol (e.g. HTTP absolute-form)
@@ -90,15 +92,19 @@ class ConnectionHandler:
                 if client_ssl or server_ssl:
                     self.establish_ssl(client=client_ssl, server=server_ssl)
 
+                if self.config.check_tcp(self.server_conn.address):
+                    self.log("Generic TCP mode for host: %s:%s" % self.server_conn.address(), "info")
+                    self.conntype = "tcp"
+
             # Delegate handling to the protocol handler
-            protocol_handler(self.conntype)(self).handle_messages()
+            protocol_handler(self.conntype)(self, **conn_kwargs).handle_messages()
 
             self.del_server_connection()
             self.log("clientdisconnect", "info")
             self.channel.tell("clientdisconnect", self)
 
         except ProxyError as e:
-            protocol_handler(self.conntype)(self).handle_error(e)
+            protocol_handler(self.conntype)(self, **conn_kwargs).handle_error(e)
         except Exception:
             import traceback, sys
 
@@ -119,14 +125,6 @@ class ConnectionHandler:
         self.server_conn = None
         self.sni = None
 
-    def check_ignore_address(self, address):
-        address = tcp.Address.wrap(address)
-        host = "%s:%s" % (address.host, address.port)
-        if host and any(rex.search(host) for rex in self.config.ignore):
-            return True
-        else:
-            return False
-
     def set_server_address(self, address):
         """
         Sets a new server address with the given priority.