Merge branch 'master' of github.com:mitmproxy/mitmproxy

4 files changed, 150 insertions, 64 deletions

diff --git a/libmproxy/proxy/config.py b/libmproxy/proxy/config.py
index dfde2958..3f579669 100644
--- a/libmproxy/proxy/config.py
+++ b/libmproxy/proxy/config.py
@@ -81,16 +81,27 @@ class ProxyConfig:
         self.check_tcp = HostMatcher(tcp_hosts)
         self.authenticator = authenticator
         self.cadir = os.path.expanduser(cadir)
-        self.certstore = certutils.CertStore.from_store(self.cadir, CONF_BASENAME)
+        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
-        self.openssl_method_client, self.openssl_options_client = version_to_openssl(ssl_version_client)
-        self.openssl_method_server, self.openssl_options_server = version_to_openssl(ssl_version_server)
+        self.openssl_method_client, self.openssl_options_client = version_to_openssl(
+            ssl_version_client)
+        self.openssl_method_server, self.openssl_options_server = version_to_openssl(
+            ssl_version_server)
         self.ssl_ports = ssl_ports
 
 
-sslversion_choices = ("all", "secure", "SSLv2", "SSLv3", "TLSv1", "TLSv1_1", "TLSv1_2")
+sslversion_choices = (
+    "all",
+    "secure",
+    "SSLv2",
+    "SSLv3",
+    "TLSv1",
+    "TLSv1_1",
+    "TLSv1_2")
 
 
 def version_to_openssl(version):
@@ -119,7 +130,8 @@ def process_proxy_options(parser, options):
     if options.transparent_proxy:
         c += 1
         if not platform.resolver:
-            return parser.error("Transparent mode not supported on this platform.")
+            return parser.error(
+                "Transparent mode not supported on this platform.")
         mode = "transparent"
     if options.socks_proxy:
         c += 1
@@ -133,28 +145,33 @@ def process_proxy_options(parser, options):
         mode = "upstream"
         upstream_server = options.upstream_proxy
     if c > 1:
-        return parser.error("Transparent, SOCKS5, reverse and upstream proxy mode "
-                            "are mutually exclusive.")
+        return parser.error(
+            "Transparent, SOCKS5, reverse and upstream proxy mode "
+            "are mutually exclusive.")
 
     if options.clientcerts:
         options.clientcerts = os.path.expanduser(options.clientcerts)
-        if not os.path.exists(options.clientcerts) or not os.path.isdir(options.clientcerts):
+        if not os.path.exists(
+                options.clientcerts) or not os.path.isdir(
+                options.clientcerts):
             return parser.error(
-                "Client certificate directory does not exist or is not a directory: %s" % options.clientcerts
-            )
+                "Client certificate directory does not exist or is not a directory: %s" %
+                options.clientcerts)
 
     if (options.auth_nonanonymous or options.auth_singleuser or options.auth_htpasswd):
         if options.auth_singleuser:
             if len(options.auth_singleuser.split(':')) != 2:
-                return parser.error("Invalid single-user specification. Please use the format username:password")
+                return parser.error(
+                    "Invalid single-user specification. Please use the format username:password")
             username, password = options.auth_singleuser.split(':')
             password_manager = http_auth.PassManSingleUser(username, password)
         elif options.auth_nonanonymous:
             password_manager = http_auth.PassManNonAnon()
         elif options.auth_htpasswd:
             try:
-                password_manager = http_auth.PassManHtpasswd(options.auth_htpasswd)
-            except ValueError, v:
+                password_manager = http_auth.PassManHtpasswd(
+                    options.auth_htpasswd)
+            except ValueError as v:
                 return parser.error(v.message)
         authenticator = http_auth.BasicProxyAuth(password_manager, "mitmproxy")
     else:
@@ -203,15 +220,18 @@ def process_proxy_options(parser, options):
 def ssl_option_group(parser):
     group = parser.add_argument_group("SSL")
     group.add_argument(
-        "--cert", dest='certs', default=[], type=str,
-        metavar="SPEC", action="append",
+        "--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. '
-             'The PEM file should contain the full certificate chain, with the leaf certificate as the first entry. '
-             'Can be passed multiple times.'
-    )
+        '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(
         "--cert-forward", action="store_true",
         dest="certforward", default=False,
@@ -238,11 +258,15 @@ 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=list(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)
-    )
+        "Defaults to %s." %
+        str(TRANSPARENT_SSL_PORTS))
     group.add_argument(
         "--ssl-version-client", dest="ssl_version_client",
         default="secure", action="store",
diff --git a/libmproxy/proxy/connection.py b/libmproxy/proxy/connection.py
index 1eeae16f..5219023b 100644
--- a/libmproxy/proxy/connection.py
+++ b/libmproxy/proxy/connection.py
@@ -7,7 +7,9 @@ from .. import stateobject, utils
 
 class ClientConnection(tcp.BaseHandler, stateobject.StateObject):
     def __init__(self, client_connection, address, server):
-        if client_connection:  # Eventually, this object is restored from state. We don't have a connection then.
+        # Eventually, this object is restored from state. We don't have a
+        # connection then.
+        if client_connection:
             tcp.BaseHandler.__init__(self, client_connection, address, server)
         else:
             self.connection = None
@@ -39,15 +41,18 @@ class ClientConnection(tcp.BaseHandler, stateobject.StateObject):
     def get_state(self, short=False):
         d = super(ClientConnection, self).get_state(short)
         d.update(
-            address={"address": self.address(), "use_ipv6": self.address.use_ipv6},
-            clientcert=self.cert.to_pem() if self.clientcert else None
-        )
+            address={
+                "address": self.address(),
+                "use_ipv6": self.address.use_ipv6},
+            clientcert=self.cert.to_pem() if self.clientcert else None)
         return d
 
     def load_state(self, state):
         super(ClientConnection, self).load_state(state)
-        self.address = tcp.Address(**state["address"]) if state["address"] else None
-        self.clientcert = certutils.SSLCert.from_pem(state["clientcert"]) if state["clientcert"] else None
+        self.address = tcp.Address(
+            **state["address"]) if state["address"] else None
+        self.clientcert = certutils.SSLCert.from_pem(
+            state["clientcert"]) if state["clientcert"] else None
 
     def copy(self):
         return copy.copy(self)
@@ -114,7 +119,7 @@ class ServerConnection(tcp.TCPClient, stateobject.StateObject):
             address={"address": self.address(),
                      "use_ipv6": self.address.use_ipv6},
             source_address= ({"address": self.source_address(),
-                             "use_ipv6": self.source_address.use_ipv6} if self.source_address else None),
+                              "use_ipv6": self.source_address.use_ipv6} if self.source_address else None),
             cert=self.cert.to_pem() if self.cert else None
         )
         return d
@@ -122,9 +127,12 @@ class ServerConnection(tcp.TCPClient, stateobject.StateObject):
     def load_state(self, state):
         super(ServerConnection, self).load_state(state)
 
-        self.address = tcp.Address(**state["address"]) if state["address"] else None
-        self.source_address = tcp.Address(**state["source_address"]) if state["source_address"] else None
-        self.cert = certutils.SSLCert.from_pem(state["cert"]) if state["cert"] else None
+        self.address = tcp.Address(
+            **state["address"]) if state["address"] else None
+        self.source_address = tcp.Address(
+            **state["source_address"]) if state["source_address"] else None
+        self.cert = certutils.SSLCert.from_pem(
+            state["cert"]) if state["cert"] else None
 
     @classmethod
     def from_state(cls, state):
@@ -147,7 +155,9 @@ class ServerConnection(tcp.TCPClient, stateobject.StateObject):
     def establish_ssl(self, clientcerts, sni, **kwargs):
         clientcert = None
         if clientcerts:
-            path = os.path.join(clientcerts, self.address.host.encode("idna")) + ".pem"
+            path = os.path.join(
+                clientcerts,
+                self.address.host.encode("idna")) + ".pem"
             if os.path.exists(path):
                 clientcert = path
         self.convert_to_ssl(cert=clientcert, sni=sni, **kwargs)
diff --git a/libmproxy/proxy/primitives.py b/libmproxy/proxy/primitives.py
index c0ae424d..9e7dae9a 100644
--- a/libmproxy/proxy/primitives.py
+++ b/libmproxy/proxy/primitives.py
@@ -1,6 +1,7 @@
 from __future__ import absolute_import
 from netlib import socks
 
+
 class ProxyError(Exception):
     def __init__(self, code, message, headers=None):
         super(ProxyError, self).__init__(message)
@@ -61,7 +62,7 @@ class TransparentProxyMode(ProxyMode):
     def get_upstream_server(self, client_conn):
         try:
             dst = self.resolver.original_addr(client_conn.connection)
-        except Exception, e:
+        except Exception as e:
             raise ProxyError(502, "Transparent mode failure: %s" % str(e))
 
         if dst[1] in self.sslports:
@@ -87,7 +88,9 @@ class Socks5ProxyMode(ProxyMode):
                 guess = ""
             raise socks.SocksError(
                 socks.REP.GENERAL_SOCKS_SERVER_FAILURE,
-                guess + "Invalid SOCKS version. Expected 0x05, got 0x%x" % msg.ver)
+                guess +
+                "Invalid SOCKS version. Expected 0x05, got 0x%x" %
+                msg.ver)
 
     def get_upstream_server(self, client_conn):
         try:
@@ -117,13 +120,15 @@ class Socks5ProxyMode(ProxyMode):
                     "mitmproxy only supports SOCKS5 CONNECT."
                 )
 
-            # We do not connect here yet, as the clientconnect event has not been handled yet.
+            # 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.
+                # dummy value, we don't have an upstream connection yet.
+                client_conn.address
             )
             connect_reply.to_file(client_conn.wfile)
             client_conn.wfile.flush()
@@ -161,4 +166,4 @@ class UpstreamProxyMode(_ConstDestinationProxyMode):
 class Log:
     def __init__(self, msg, level="info"):
         self.msg = msg
-        self.level = level
\ No newline at end of file
+        self.level = level
diff --git a/libmproxy/proxy/server.py b/libmproxy/proxy/server.py
index a72f9aba..e1587df1 100644
--- a/libmproxy/proxy/server.py
+++ b/libmproxy/proxy/server.py
@@ -34,7 +34,7 @@ class ProxyServer(tcp.TCPServer):
         self.config = config
         try:
             tcp.TCPServer.__init__(self, (config.host, config.port))
-        except socket.error, v:
+        except socket.error as v:
             raise ProxyServerError('Error starting proxy server: ' + repr(v))
         self.channel = None
 
@@ -46,16 +46,30 @@ class ProxyServer(tcp.TCPServer):
         self.channel = channel
 
     def handle_client_connection(self, conn, client_address):
-        h = ConnectionHandler(self.config, conn, client_address, self, self.channel)
+        h = ConnectionHandler(
+            self.config,
+            conn,
+            client_address,
+            self,
+            self.channel)
         h.handle()
         h.finish()
 
 
 class ConnectionHandler:
-    def __init__(self, config, client_connection, client_address, server, channel):
+    def __init__(
+            self,
+            config,
+            client_connection,
+            client_address,
+            server,
+            channel):
         self.config = config
         """@type: libmproxy.proxy.config.ProxyConfig"""
-        self.client_conn = ClientConnection(client_connection, client_address, server)
+        self.client_conn = ClientConnection(
+            client_connection,
+            client_address,
+            server)
         """@type: libmproxy.proxy.connection.ClientConnection"""
         self.server_conn = None
         """@type: libmproxy.proxy.connection.ServerConnection"""
@@ -70,17 +84,23 @@ 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)
+            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.config.check_ignore(self.server_conn.address):
-                    self.log("Ignore host: %s:%s" % self.server_conn.address(), "info")
+                    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)
+                # No upstream info from the metadata: upstream info in the
+                # protocol (e.g. HTTP absolute-form)
+                pass
 
             self.channel.ask("clientconnect", self)
 
@@ -92,11 +112,17 @@ class ConnectionHandler:
                     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.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, **conn_kwargs).handle_messages()
+            protocol_handler(
+                self.conntype)(
+                self,
+                **conn_kwargs).handle_messages()
 
             self.log("clientdisconnect", "info")
             self.channel.tell("clientdisconnect", self)
@@ -104,7 +130,8 @@ class ConnectionHandler:
         except ProxyError as e:
             protocol_handler(self.conntype)(self, **conn_kwargs).handle_error(e)
         except Exception:
-            import traceback, sys
+            import traceback
+            import sys
 
             self.log(traceback.format_exc(), "error")
             print >> sys.stderr, traceback.format_exc()
@@ -112,7 +139,8 @@ class ConnectionHandler:
             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.
+            # The client connection is closed by the ProxyServer and does not
+            # have be handled here.
             self.del_server_connection()
 
     def del_server_connection(self):
@@ -122,8 +150,10 @@ 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.log(
+                "serverdisconnect", "debug", [
+                    "%s:%s" %
+                    (self.server_conn.address.host, self.server_conn.address.port)])
             self.channel.tell("serverdisconnect", self)
         self.server_conn = None
 
@@ -141,7 +171,9 @@ class ConnectionHandler:
         if self.server_conn:
             self.del_server_connection()
 
-        self.log("Set new server address: %s:%s" % (address.host, address.port), "debug")
+        self.log(
+            "Set new server address: %s:%s" %
+            (address.host, address.port), "debug")
         self.server_conn = ServerConnection(address)
 
     def establish_server_connection(self, ask=True):
@@ -155,12 +187,16 @@ class ConnectionHandler:
         """
         if self.server_conn.connection:
             return
-        self.log("serverconnect", "debug", ["%s:%s" % self.server_conn.address()[:2]])
+        self.log(
+            "serverconnect", "debug", [
+                "%s:%s" %
+                self.server_conn.address()[
+                    :2]])
         if ask:
             self.channel.ask("serverconnect", self)
         try:
             self.server_conn.connect()
-        except tcp.NetLibError, v:
+        except tcp.NetLibError as v:
             raise ProxyError(502, v)
 
     def establish_ssl(self, client=False, server=False, sni=None):
@@ -237,7 +273,8 @@ class ConnectionHandler:
         self.server_conn.state = state
 
         # Receiving new_sni where had_ssl is False is a weird case that happens when the workaround for
-        # https://github.com/mitmproxy/mitmproxy/issues/427 is active. In this case, we want to establish SSL as well.
+        # https://github.com/mitmproxy/mitmproxy/issues/427 is active. In this
+        # case, we want to establish SSL as well.
         if had_ssl or new_sni:
             self.establish_ssl(server=True, sni=sni)
 
@@ -246,8 +283,10 @@ class ConnectionHandler:
 
     def log(self, msg, level, subs=()):
         msg = [
-            "%s:%s: %s" % (self.client_conn.address.host, self.client_conn.address.port, msg)
-        ]
+            "%s:%s: %s" %
+            (self.client_conn.address.host,
+             self.client_conn.address.port,
+             msg)]
         for i in subs:
             msg.append("  -> " + i)
         msg = "\n".join(msg)
@@ -255,11 +294,13 @@ 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), None
+            return self.server_conn.cert, self.config.certstore.gen_pkey(
+                self.server_conn.cert), None
         else:
             host = self.server_conn.address.host
             sans = []
-            if self.server_conn.ssl_established and (not self.config.no_upstream_cert):
+            if self.server_conn.ssl_established and (
+                    not self.config.no_upstream_cert):
                 upstream_cert = self.server_conn.cert
                 sans.extend(upstream_cert.altnames)
                 if upstream_cert.cn:
@@ -291,8 +332,11 @@ class ConnectionHandler:
                 #   - We established SSL with the server previously
                 #   - We initially wanted to establish SSL with the server,
                 #     but the server refused to negotiate without SNI.
-                if self.server_conn.ssl_established or hasattr(self.server_conn, "may_require_sni"):
-                    self.server_reconnect(sni)  # reconnect to upstream server with SNI
+                if self.server_conn.ssl_established or hasattr(
+                        self.server_conn,
+                        "may_require_sni"):
+                    # reconnect to upstream server with SNI
+                    self.server_reconnect(sni)
                 # Now, change client context to reflect changed certificate:
                 cert, key, chain_file = self.find_cert()
                 new_context = self.client_conn.create_ssl_context(
@@ -308,4 +352,7 @@ class ConnectionHandler:
         # make dang sure it doesn't happen.
         except:  # pragma: no cover
             import traceback
-            self.log("Error in handle_sni:\r\n" + traceback.format_exc(), "error")
+            self.log(
+                "Error in handle_sni:\r\n" +
+                traceback.format_exc(),
+                "error")