Merge remote-tracking branch 'base/master'

3 files changed, 164 insertions, 50 deletions

diff --git a/libmproxy/proxy/config.py b/libmproxy/proxy/config.py
index 62104a24..3d373a28 100644
--- a/libmproxy/proxy/config.py
+++ b/libmproxy/proxy/config.py
@@ -1,26 +1,54 @@
 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
+from .primitives import RegularProxyMode, TransparentProxyMode, UpstreamProxyMode, ReverseProxyMode, Socks5ProxyMode
 
 TRANSPARENT_SSL_PORTS = [443, 8443]
 CONF_BASENAME = "mitmproxy"
-CONF_DIR = "~/.mitmproxy"
+CA_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:
-    def __init__(self, host='', port=8080, server_version=version.NAMEVERSION,
-                 confdir=CONF_DIR, ca_file=None, 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=[],
-                 ciphers=None, certs=[], certforward=False, ssl_ports=TRANSPARENT_SSL_PORTS):
+    def __init__(
+            self,
+            host='',
+            port=8080,
+            server_version=version.NAMEVERSION,
+            cadir=CA_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_hosts=[],
+            tcp_hosts=[],
+            ciphers=None,
+            certs=[],
+            certforward=False,
+            ssl_ports=TRANSPARENT_SSL_PORTS
+    ):
         self.host = host
         self.port = port
         self.server_version = server_version
@@ -30,7 +58,9 @@ class ProxyConfig:
         self.body_size_limit = body_size_limit
 
         if mode == "transparent":
-            self.mode = TransparentProxyMode(platform.resolver(), TRANSPARENT_SSL_PORTS)
+            self.mode = TransparentProxyMode(platform.resolver(), ssl_ports)
+        elif mode == "socks5":
+            self.mode = Socks5ProxyMode(ssl_ports)
         elif mode == "reverse":
             self.mode = ReverseProxyMode(upstream_server)
         elif mode == "upstream":
@@ -42,11 +72,11 @@ 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.ca_file = ca_file or os.path.join(self.confdir, CONF_BASENAME + "-ca.pem")
-        self.certstore = certutils.CertStore.from_store(self.confdir, CONF_BASENAME)
+        self.cadir = os.path.expanduser(cadir)
+        self.certstore = certutils.CertStore.from_store(self.cadir, CONF_BASENAME)
         for spec, cert in certs:
             self.certstore.add_cert_file(spec, cert)
         self.certforward = certforward
@@ -63,6 +93,9 @@ def process_proxy_options(parser, options):
         if not platform.resolver:
             return parser.error("Transparent mode not supported on this platform.")
         mode = "transparent"
+    if options.socks_proxy:
+        c += 1
+        mode = "socks5"
     if options.reverse_proxy:
         c += 1
         mode = "reverse"
@@ -72,7 +105,7 @@ def process_proxy_options(parser, options):
         mode = "upstream"
         upstream_server = options.upstream_proxy
     if c > 1:
-        return parser.error("Transparent mode, reverse mode and upstream proxy mode "
+        return parser.error("Transparent, SOCKS5, reverse and upstream proxy mode "
                             "are mutually exclusive.")
 
     if options.clientcerts:
@@ -109,10 +142,16 @@ def process_proxy_options(parser, options):
             parser.error("Certificate file does not exist: %s" % parts[1])
         certs.append(parts)
 
+    ssl_ports = options.ssl_ports
+    if options.ssl_ports != TRANSPARENT_SSL_PORTS:
+        # arparse appends to default value by default, strip that off.
+        # see http://bugs.python.org/issue16399
+        ssl_ports = ssl_ports[len(TRANSPARENT_SSL_PORTS):]
+
     return ProxyConfig(
         host=options.addr,
         port=options.port,
-        confdir=options.confdir,
+        cadir=options.cadir,
         clientcerts=options.clientcerts,
         no_upstream_cert=options.no_upstream_cert,
         body_size_limit=body_size_limit,
@@ -120,11 +159,13 @@ 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,
         certforward=options.certforward,
+        ssl_ports=ssl_ports
     )
 
 
@@ -133,10 +174,12 @@ def ssl_option_group(parser):
     group.add_argument(
         "--cert", dest='certs', default=[], type=str,
         metavar="SPEC", action="append",
-        help='Add an SSL certificate. SPEC is of the form "[domain=]path". ' \
-             'The domain may include a wildcard, and is equal to "*" if not specified. ' \
-             'The file at path is a certificate in PEM format. If a private key is included in the PEM, ' \
-             'it is used, else the default key in the conf dir is used. Can be passed multiple times.'
+        help='Add an SSL certificate. SPEC is of the form "[domain=]path". '
+             'The domain may include a wildcard, and is equal to "*" if not specified. '
+             'The file at path is a certificate in PEM format. If a private key is included in the PEM, '
+             'it is used, else the default key in the conf dir is used. '
+             'The PEM file should contain the full certificate chain, with the leaf certificate as the first entry. '
+             'Can be passed multiple times.'
     )
     group.add_argument(
         "--client-certs", action="store",
@@ -159,7 +202,7 @@ def ssl_option_group(parser):
         help="Don't connect to upstream server to look up certificate details."
     )
     group.add_argument(
-        "--ssl-port", action="append", type=int, dest="ssl_ports", default=TRANSPARENT_SSL_PORTS,
+        "--ssl-port", action="append", type=int, dest="ssl_ports", default=list(TRANSPARENT_SSL_PORTS),
         metavar="PORT",
         help="Can be passed multiple times. Specify destination ports which are assumed to be SSL. "
              "Defaults to %s." % str(TRANSPARENT_SSL_PORTS)
diff --git a/libmproxy/proxy/primitives.py b/libmproxy/proxy/primitives.py
index 23d089d3..c0ae424d 100644
--- a/libmproxy/proxy/primitives.py
+++ b/libmproxy/proxy/primitives.py
@@ -1,5 +1,5 @@
 from __future__ import absolute_import
-
+from netlib import socks
 
 class ProxyError(Exception):
     def __init__(self, code, message, headers=None):
@@ -15,7 +15,7 @@ class ProxyMode(object):
     http_form_in = None
     http_form_out = None
 
-    def get_upstream_server(self, conn):
+    def get_upstream_server(self, client_conn):
         """
         Returns the address of the server to connect to.
         Returns None if the address needs to be determined on the protocol level (regular proxy mode)
@@ -46,7 +46,7 @@ class RegularProxyMode(ProxyMode):
     http_form_in = "absolute"
     http_form_out = "relative"
 
-    def get_upstream_server(self, conn):
+    def get_upstream_server(self, client_conn):
         return None
 
 
@@ -58,9 +58,9 @@ class TransparentProxyMode(ProxyMode):
         self.resolver = resolver
         self.sslports = sslports
 
-    def get_upstream_server(self, conn):
+    def get_upstream_server(self, client_conn):
         try:
-            dst = self.resolver.original_addr(conn)
+            dst = self.resolver.original_addr(client_conn.connection)
         except Exception, e:
             raise ProxyError(502, "Transparent mode failure: %s" % str(e))
 
@@ -71,11 +71,80 @@ class TransparentProxyMode(ProxyMode):
         return [ssl, ssl] + list(dst)
 
 
+class Socks5ProxyMode(ProxyMode):
+    http_form_in = "relative"
+    http_form_out = "relative"
+
+    def __init__(self, sslports):
+        self.sslports = sslports
+
+    @staticmethod
+    def _assert_socks5(msg):
+        if msg.ver != socks.VERSION.SOCKS5:
+            if msg.ver == ord("G") and len(msg.methods) == ord("E"):
+                guess = "Probably not a SOCKS request but a regular HTTP request. "
+            else:
+                guess = ""
+            raise socks.SocksError(
+                socks.REP.GENERAL_SOCKS_SERVER_FAILURE,
+                guess + "Invalid SOCKS version. Expected 0x05, got 0x%x" % msg.ver)
+
+    def get_upstream_server(self, client_conn):
+        try:
+            # Parse Client Greeting
+            client_greet = socks.ClientGreeting.from_file(client_conn.rfile)
+            self._assert_socks5(client_greet)
+            if socks.METHOD.NO_AUTHENTICATION_REQUIRED not in client_greet.methods:
+                raise socks.SocksError(
+                    socks.METHOD.NO_ACCEPTABLE_METHODS,
+                    "mitmproxy only supports SOCKS without authentication"
+                )
+
+            # Send Server Greeting
+            server_greet = socks.ServerGreeting(
+                socks.VERSION.SOCKS5,
+                socks.METHOD.NO_AUTHENTICATION_REQUIRED
+            )
+            server_greet.to_file(client_conn.wfile)
+            client_conn.wfile.flush()
+
+            # Parse Connect Request
+            connect_request = socks.Message.from_file(client_conn.rfile)
+            self._assert_socks5(connect_request)
+            if connect_request.msg != socks.CMD.CONNECT:
+                raise socks.SocksError(
+                    socks.REP.COMMAND_NOT_SUPPORTED,
+                    "mitmproxy only supports SOCKS5 CONNECT."
+                )
+
+            # We do not connect here yet, as the clientconnect event has not been handled yet.
+
+            connect_reply = socks.Message(
+                socks.VERSION.SOCKS5,
+                socks.REP.SUCCEEDED,
+                socks.ATYP.DOMAINNAME,
+                client_conn.address  # dummy value, we don't have an upstream connection yet.
+            )
+            connect_reply.to_file(client_conn.wfile)
+            client_conn.wfile.flush()
+
+            ssl = bool(connect_request.addr.port in self.sslports)
+            return ssl, ssl, connect_request.addr.host, connect_request.addr.port
+
+        except socks.SocksError as e:
+            msg = socks.Message(5, e.code, socks.ATYP.DOMAINNAME, repr(e))
+            try:
+                msg.to_file(client_conn.wfile)
+            except:
+                pass
+            raise ProxyError(502, "SOCKS5 mode failure: %s" % str(e))
+
+
 class _ConstDestinationProxyMode(ProxyMode):
     def __init__(self, dst):
         self.dst = dst
 
-    def get_upstream_server(self, conn):
+    def get_upstream_server(self, client_conn):
         return self.dst
 
 
diff --git a/libmproxy/proxy/server.py b/libmproxy/proxy/server.py
index 307a4bcd..55e2b30e 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
-            upstream_info = self.config.mode.get_upstream_server(self.client_conn.connection)
+            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,18 @@ 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
 
@@ -106,6 +111,10 @@ class ConnectionHandler:
             print >> sys.stderr, traceback.format_exc()
             print >> sys.stderr, "mitmproxy has crashed!"
             print >> sys.stderr, "Please lodge a bug report at: https://github.com/mitmproxy/mitmproxy"
+        finally:
+            # Make sure that we close the server connection in any case.
+            # The client connection is closed by the ProxyServer and does not have be handled here.
+            self.del_server_connection()
 
     def del_server_connection(self):
         """
@@ -113,20 +122,13 @@ class ConnectionHandler:
         """
         if self.server_conn and self.server_conn.connection:
             self.server_conn.finish()
+            self.server_conn.close()
             self.log("serverdisconnect", "debug", ["%s:%s" % (self.server_conn.address.host,
                                                               self.server_conn.address.port)])
             self.channel.tell("serverdisconnect", self)
         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.
@@ -190,14 +192,14 @@ class ConnectionHandler:
         if client:
             if self.client_conn.ssl_established:
                 raise ProxyError(502, "SSL to Client already established.")
-            cert, key = self.find_cert()
+            cert, key, chain_file = self.find_cert()
             try:
                 self.client_conn.convert_to_ssl(
                     cert, key,
                     handle_sni=self.handle_sni,
                     cipher_list=self.config.ciphers,
                     dhparams=self.config.certstore.dhparams,
-                    ca_file=self.config.ca_file
+                    chain_file=chain_file
                 )
             except tcp.NetLibError as v:
                 raise ProxyError(400, repr(v))
@@ -234,7 +236,7 @@ class ConnectionHandler:
 
     def find_cert(self):
         if self.config.certforward and self.server_conn.ssl_established:
-            return self.server_conn.cert, self.config.certstore.gen_pkey(self.server_conn.cert)
+            return self.server_conn.cert, self.config.certstore.gen_pkey(self.server_conn.cert), None
         else:
             host = self.server_conn.address.host
             sans = []
@@ -264,17 +266,17 @@ class ConnectionHandler:
                 self.log("SNI received: %s" % self.sni, "debug")
                 self.server_reconnect()  # reconnect to upstream server with SNI
                 # Now, change client context to reflect changed certificate:
-                cert, key = self.find_cert()
+                cert, key, chain_file = self.find_cert()
                 new_context = self.client_conn._create_ssl_context(
                     cert, key,
                     method=SSL.TLSv1_METHOD,
                     cipher_list=self.config.ciphers,
                     dhparams=self.config.certstore.dhparams,
-                    ca_file=self.config.ca_file
+                    chain_file=chain_file
                 )
                 connection.set_context(new_context)
         # An unhandled exception in this method will core dump PyOpenSSL, so
         # make dang sure it doesn't happen.
-        except Exception:  # pragma: no cover
+        except:  # pragma: no cover
             import traceback
-            self.log("Error in handle_sni:\r\n" + traceback.format_exc(), "error")
\ No newline at end of file
+            self.log("Error in handle_sni:\r\n" + traceback.format_exc(), "error")