restructure code, remove cruft

4 files changed, 51 insertions, 548 deletions

diff --git a/libmproxy/proxy/__init__.py b/libmproxy/proxy/__init__.py
index 7d664707..709654cb 100644
--- a/libmproxy/proxy/__init__.py
+++ b/libmproxy/proxy/__init__.py
@@ -1,2 +1,11 @@
-from .primitives import *
-from .config import ProxyConfig, process_proxy_options
\ No newline at end of file
+from __future__ import (absolute_import, print_function, division)
+
+from .primitives import Log, Kill
+from .config import ProxyConfig
+from .connection import ClientConnection, ServerConnection
+
+__all__ = [
+    "Log", "Kill",
+    "ProxyConfig",
+    "ClientConnection", "ServerConnection"
+]
\ No newline at end of file
diff --git a/libmproxy/proxy/connection.py b/libmproxy/proxy/connection.py
index c329ed64..94f318f6 100644
--- a/libmproxy/proxy/connection.py
+++ b/libmproxy/proxy/connection.py
@@ -12,7 +12,7 @@ class ClientConnection(tcp.BaseHandler, stateobject.StateObject):
         # 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)
+            super(ClientConnection, self).__init__(client_connection, address, server)
         else:
             self.connection = None
             self.server = None
@@ -80,11 +80,11 @@ class ClientConnection(tcp.BaseHandler, stateobject.StateObject):
         return f
 
     def convert_to_ssl(self, *args, **kwargs):
-        tcp.BaseHandler.convert_to_ssl(self, *args, **kwargs)
+        super(ClientConnection, self).convert_to_ssl(*args, **kwargs)
         self.timestamp_ssl_setup = utils.timestamp()
 
     def finish(self):
-        tcp.BaseHandler.finish(self)
+        super(ClientConnection, self).finish()
         self.timestamp_end = utils.timestamp()
 
 
diff --git a/libmproxy/proxy/primitives.py b/libmproxy/proxy/primitives.py
index a9f31181..2e440fe8 100644
--- a/libmproxy/proxy/primitives.py
+++ b/libmproxy/proxy/primitives.py
@@ -1,178 +1,15 @@
 from __future__ import absolute_import
+import collections
 from netlib import socks, tcp
 
 
-class ProxyError(Exception):
-    def __init__(self, code, message, headers=None):
-        super(ProxyError, self).__init__(message)
-        self.code, self.headers = code, headers
-
-
-class ProxyServerError(Exception):
-    pass
-
-
-class ProxyMode(object):
-    http_form_in = None
-    http_form_out = None
-
-    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)
-        """
-        raise NotImplementedError()  # pragma: nocover
-
-    @property
-    def name(self):
-        return self.__class__.__name__.replace("ProxyMode", "").lower()
-
-    def __str__(self):
-        return self.name
-
-    def __eq__(self, other):
-        """
-        Allow comparisions with "regular" etc.
-        """
-        if isinstance(other, ProxyMode):
-            return self is other
-        else:
-            return self.name == other
-
-    def __ne__(self, other):
-        return not self.__eq__(other)
-
-
-class RegularProxyMode(ProxyMode):
-    http_form_in = "absolute"
-    http_form_out = "relative"
-
-    def get_upstream_server(self, client_conn):
-        return None
-
-
-class SpoofMode(ProxyMode):
-    http_form_in = "relative"
-    http_form_out = "relative"
-
-    def get_upstream_server(self, client_conn):
-        return None
-
-    @property
-    def name(self):
-        return "spoof"
-
-
-class SSLSpoofMode(ProxyMode):
-    http_form_in = "relative"
-    http_form_out = "relative"
-
-    def __init__(self, sslport):
-        self.sslport = sslport
-
-    def get_upstream_server(self, client_conn):
-        return None
-
-    @property
-    def name(self):
-        return "sslspoof"
-
-
-class TransparentProxyMode(ProxyMode):
-    http_form_in = "relative"
-    http_form_out = "relative"
-
-    def __init__(self, resolver, sslports):
-        self.resolver = resolver
-        self.sslports = sslports
-
-    def get_upstream_server(self, client_conn):
-        try:
-            dst = self.resolver.original_addr(client_conn.connection)
-        except Exception as e:
-            raise ProxyError(502, "Transparent mode failure: %s" % str(e))
-
-        if dst[1] in self.sslports:
-            ssl = True
-        else:
-            ssl = False
-        return [ssl, ssl] + list(dst)
-
-
-class Socks5ProxyMode(ProxyMode):
-    http_form_in = "relative"
-    http_form_out = "relative"
-
-    def __init__(self, sslports):
-        self.sslports = sslports
-
-    def get_upstream_server(self, client_conn):
-        try:
-            # Parse Client Greeting
-            client_greet = socks.ClientGreeting.from_file(client_conn.rfile, fail_early=True)
-            client_greet.assert_socks5()
-            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)
-            connect_request.assert_socks5()
-            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,
-                connect_request.atyp,
-                # dummy value, we don't have an upstream connection yet.
-                connect_request.addr
-            )
-            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, tcp.NetLibError) as e:
-            raise ProxyError(502, "SOCKS5 mode failure: %s" % repr(e))
-
-
-class _ConstDestinationProxyMode(ProxyMode):
-    def __init__(self, dst):
-        self.dst = dst
-
-    def get_upstream_server(self, client_conn):
-        return self.dst
-
-
-class ReverseProxyMode(_ConstDestinationProxyMode):
-    http_form_in = "relative"
-    http_form_out = "relative"
-
-
-class UpstreamProxyMode(_ConstDestinationProxyMode):
-    http_form_in = "absolute"
-    http_form_out = "absolute"
-
-
-class Log:
+class Log(object):
     def __init__(self, msg, level="info"):
         self.msg = msg
         self.level = level
+
+
+class Kill(Exception):
+    """
+    Kill a connection.
+    """
\ No newline at end of file
diff --git a/libmproxy/proxy/server.py b/libmproxy/proxy/server.py
index 1fc4cbda..69784014 100644
--- a/libmproxy/proxy/server.py
+++ b/libmproxy/proxy/server.py
@@ -3,16 +3,14 @@ from __future__ import absolute_import, print_function
 import traceback
 import sys
 import socket
-from libmproxy.protocol2.layer import Kill
 from netlib import tcp
 from netlib.http.http1 import HTTP1Protocol
 from netlib.tcp import NetLibError
 
-from ..protocol.handle import protocol_handler
 from .. import protocol2
-from ..exceptions import ProtocolException
-from .primitives import ProxyServerError, Log, ProxyError
-from .connection import ClientConnection, ServerConnection
+from ..exceptions import ProtocolException, ServerException
+from .primitives import Log, Kill
+from .connection import ClientConnection
 
 
 class DummyServer:
@@ -38,9 +36,9 @@ class ProxyServer(tcp.TCPServer):
         """
         self.config = config
         try:
-            tcp.TCPServer.__init__(self, (config.host, config.port))
-        except socket.error as v:
-            raise ProxyServerError('Error starting proxy server: ' + repr(v))
+            super(ProxyServer, self).__init__((config.host, config.port))
+        except socket.error as e:
+            raise ServerException('Error starting proxy server: ' + repr(e), e)
         self.channel = None
 
     def start_slave(self, klass, channel):
@@ -51,33 +49,28 @@ class ProxyServer(tcp.TCPServer):
         self.channel = channel
 
     def handle_client_connection(self, conn, client_address):
-        h = ConnectionHandler2(
-            self.config,
+        h = ConnectionHandler(
             conn,
             client_address,
-            self,
-            self.channel)
+            self.config,
+            self.channel
+        )
         h.handle()
-        h.finish()
 
 
-class ConnectionHandler2:
-    # FIXME: parameter ordering
-    # FIXME: remove server attribute
-    def __init__(self, config, client_conn, client_address, server, channel):
+class ConnectionHandler(object):
+    def __init__(self, client_conn, client_address, config, channel):
         self.config = config
         """@type: libmproxy.proxy.config.ProxyConfig"""
         self.client_conn = ClientConnection(
             client_conn,
             client_address,
-            server)
+            None)
         """@type: libmproxy.proxy.connection.ClientConnection"""
         self.channel = channel
         """@type: libmproxy.controller.Channel"""
 
-    def handle(self):
-        self.log("clientconnect", "info")
-
+    def _create_root_layer(self):
         root_context = protocol2.RootContext(
             self.client_conn,
             self.config,
@@ -86,33 +79,38 @@ class ConnectionHandler2:
 
         mode = self.config.mode
         if mode == "upstream":
-            root_layer = protocol2.HttpUpstreamProxy(
+            return protocol2.HttpUpstreamProxy(
                 root_context,
                 self.config.upstream_server.address
             )
         elif mode == "transparent":
-            root_layer = protocol2.TransparentProxy(root_context)
+            return protocol2.TransparentProxy(root_context)
         elif mode == "reverse":
             client_tls = self.config.upstream_server.scheme.startswith("https")
             server_tls = self.config.upstream_server.scheme.endswith("https")
-            root_layer = protocol2.ReverseProxy(
+            return protocol2.ReverseProxy(
                 root_context,
                 self.config.upstream_server.address,
                 client_tls,
                 server_tls
             )
         elif mode == "socks5":
-            root_layer = protocol2.Socks5Proxy(root_context)
+            return protocol2.Socks5Proxy(root_context)
         elif mode == "regular":
-            root_layer = protocol2.HttpProxy(root_context)
+            return protocol2.HttpProxy(root_context)
         elif callable(mode):  # pragma: nocover
-            root_layer = mode(root_context)
+            return mode(root_context)
         else:  # pragma: nocover
             raise ValueError("Unknown proxy mode: %s" % mode)
 
+    def handle(self):
+        self.log("clientconnect", "info")
+
+        root_layer = self._create_root_layer()
+
         try:
             root_layer()
-        except Kill as e:
+        except Kill:
             self.log("Connection killed", "info")
         except ProtocolException as e:
             self.log(e, "info")
@@ -131,349 +129,8 @@ class ConnectionHandler2:
             print("Please lodge a bug report at: https://github.com/mitmproxy/mitmproxy", file=sys.stderr)
 
         self.log("clientdisconnect", "info")
-
-    def finish(self):
-        self.client_conn.finish()
-
-    def log(self, msg, level, subs=()):
-        # FIXME: Duplicate code
-        full_msg = [
-            "%s:%s: %s" %
-            (self.client_conn.address.host,
-             self.client_conn.address.port,
-             msg)]
-        for i in subs:
-            full_msg.append("  -> " + i)
-        full_msg = "\n".join(full_msg)
-        self.channel.tell("log", Log(full_msg, level))
-
-
-class ConnectionHandler:
-    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)
-        """@type: libmproxy.proxy.connection.ClientConnection"""
-        self.server_conn = None
-        """@type: libmproxy.proxy.connection.ServerConnection"""
-        self.channel = channel
-        """@type: libmproxy.controller.Channel"""
-
-        self.conntype = "http"
-
-    def handle(self):
-        try:
-            self.log("clientconnect", "info")
-
-            # 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.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:
-                # No upstream info from the metadata: upstream info in the
-                # protocol (e.g. HTTP absolute-form)
-                pass
-
-            self.channel.ask("clientconnect", self)
-
-            # Check for existing connection: If an inline script already established a
-            # connection, do not apply client_ssl or server_ssl.
-            if self.server_conn and not self.server_conn.connection:
-                self.establish_server_connection()
-                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"
-
-            elif not self.server_conn and self.config.mode == "sslspoof":
-                port = self.config.mode.sslport
-                self.set_server_address(("-", port))
-                self.establish_ssl(client=True)
-                host = self.client_conn.connection.get_servername()
-                if host:
-                    self.set_server_address((host, port))
-                    self.establish_server_connection()
-                    self.establish_ssl(server=True, sni=host)
-
-            # Delegate handling to the protocol handler
-            protocol_handler(
-                self.conntype)(
-                self,
-                **conn_kwargs).handle_messages()
-
-            self.log("clientdisconnect", "info")
-            self.channel.tell("clientdisconnect", self)
-
-        except ProxyError as e:
-            protocol_handler(self.conntype)(self, **conn_kwargs).handle_error(e)
-        except Exception:
-            import traceback
-            import sys
-
-            self.log(traceback.format_exc(), "error")
-            print(traceback.format_exc(), file=sys.stderr)
-            print("mitmproxy has crashed!", file=sys.stderr)
-            print("Please lodge a bug report at: https://github.com/mitmproxy/mitmproxy", file=sys.stderr)
-        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):
-        """
-        Deletes (and closes) an existing server connection.
-        """
-        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
-
-    def set_server_address(self, addr):
-        """
-        Sets a new server address with the given priority.
-        Does not re-establish either connection or SSL handshake.
-        """
-        address = tcp.Address.wrap(addr)
-
-        # Don't reconnect to the same destination.
-        if self.server_conn and self.server_conn.address == address:
-            return
-
-        if self.server_conn:
-            self.del_server_connection()
-
-        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):
-        """
-        Establishes a new server connection.
-        If there is already an existing server connection, the function returns immediately.
-
-        By default, this function ".ask"s the proxy master. This is deadly if this function is already called from the
-        master (e.g. via change_server), because this navigates us in a simple deadlock (the master is single-threaded).
-        In these scenarios, ask=False can be passed to suppress the call to the master.
-        """
-        if self.server_conn.connection:
-            return
-        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 as v:
-            raise ProxyError(502, v)
-
-    def establish_ssl(self, client=False, server=False, sni=None):
-        """
-        Establishes SSL on the existing connection(s) to the server or the client,
-        as specified by the parameters.
-        """
-
-        # Logging
-        if client or server:
-            subs = []
-            if client:
-                subs.append("with client")
-            if server:
-                subs.append("with server (sni: %s)" % sni)
-            self.log("Establish SSL", "debug", subs)
-
-        if server:
-            if not self.server_conn or not self.server_conn.connection:
-                raise ProxyError(502, "No server connection.")
-            if self.server_conn.ssl_established:
-                raise ProxyError(502, "SSL to Server already established.")
-            try:
-                self.server_conn.establish_ssl(
-                    self.config.clientcerts,
-                    sni,
-                    method=self.config.openssl_method_server,
-                    options=self.config.openssl_options_server,
-                    verify_options=self.config.openssl_verification_mode_server,
-                    ca_path=self.config.openssl_trusted_cadir_server,
-                    ca_pemfile=self.config.openssl_trusted_ca_server,
-                    cipher_list=self.config.ciphers_server,
-                )
-                ssl_cert_err = self.server_conn.ssl_verification_error
-                if ssl_cert_err is not None:
-                    self.log(
-                        "SSL verification failed for upstream server at depth %s with error: %s" %
-                        (ssl_cert_err['depth'], ssl_cert_err['errno']),
-                        "error")
-                    self.log("Ignoring server verification error, continuing with connection", "error")
-            except tcp.NetLibError as v:
-                e = ProxyError(502, repr(v))
-                # Workaround for https://github.com/mitmproxy/mitmproxy/issues/427
-                # The upstream server may reject connections without SNI, which means we need to
-                # establish SSL with the client first, hope for a SNI (which triggers a reconnect which replaces the
-                # ServerConnection object) and see whether that worked.
-                if client and "handshake failure" in e.message:
-                    self.server_conn.may_require_sni = e
-                else:
-                    ssl_cert_err = self.server_conn.ssl_verification_error
-                    if ssl_cert_err is not None:
-                        self.log(
-                            "SSL verification failed for upstream server at depth %s with error: %s" %
-                            (ssl_cert_err['depth'], ssl_cert_err['errno']),
-                            "error")
-                        self.log("Aborting connection attempt", "error")
-                    raise e
-        if client:
-            if self.client_conn.ssl_established:
-                raise ProxyError(502, "SSL to Client already established.")
-            cert, key, chain_file = self.find_cert()
-            try:
-                self.client_conn.convert_to_ssl(
-                    cert, key,
-                    method=self.config.openssl_method_client,
-                    options=self.config.openssl_options_client,
-                    handle_sni=self.handle_sni,
-                    cipher_list=self.config.ciphers_client,
-                    dhparams=self.config.certstore.dhparams,
-                    chain_file=chain_file
-                )
-            except tcp.NetLibError as v:
-                raise ProxyError(400, repr(v))
-
-            # Workaround for #427 part 2
-            if server and hasattr(self.server_conn, "may_require_sni"):
-                raise self.server_conn.may_require_sni
-
-    def server_reconnect(self, new_sni=False):
-        address = self.server_conn.address
-        had_ssl = self.server_conn.ssl_established
-        state = self.server_conn.state
-        sni = new_sni or self.server_conn.sni
-        self.log("(server reconnect follows)", "debug")
-        self.del_server_connection()
-        self.set_server_address(address)
-        self.establish_server_connection()
-
-        for s in state:
-            protocol_handler(s[0])(self).handle_server_reconnect(s[1])
-        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.
-        if had_ssl or new_sni:
-            self.establish_ssl(server=True, sni=sni)
-
-    def finish(self):
         self.client_conn.finish()
 
-    def log(self, msg, level, subs=()):
-        full_msg = [
-            "%s:%s: %s" %
-            (self.client_conn.address.host,
-             self.client_conn.address.port,
-             msg)]
-        for i in subs:
-            full_msg.append("  -> " + i)
-        full_msg = "\n".join(full_msg)
-        self.channel.tell("log", Log(full_msg, level))
-
-    def find_cert(self):
-        host = self.server_conn.address.host
-        sans = []
-        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:
-                sans.append(host)
-                host = upstream_cert.cn.decode("utf8").encode("idna")
-        if self.server_conn.sni:
-            sans.append(self.server_conn.sni)
-        # for ssl spoof mode
-        if hasattr(self.client_conn, "sni"):
-            sans.append(self.client_conn.sni)
-
-        ret = self.config.certstore.get_cert(host, sans)
-        if not ret:
-            raise ProxyError(502, "Unable to generate dummy cert.")
-        return ret
-
-    def handle_sni(self, connection):
-        """
-        This callback gets called during the SSL handshake with the client.
-        The client has just sent the Sever Name Indication (SNI). We now connect upstream to
-        figure out which certificate needs to be served.
-        """
-        try:
-            sn = connection.get_servername()
-            if not sn:
-                return
-            sni = sn.decode("utf8").encode("idna")
-            # for ssl spoof mode
-            self.client_conn.sni = sni
-
-            if sni != self.server_conn.sni:
-                self.log("SNI received: %s" % sni, "debug")
-                # We should only re-establish upstream SSL if one of the following conditions is true:
-                #   - 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"):
-                    # 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(
-                    cert, key,
-                    method=self.config.openssl_method_client,
-                    options=self.config.openssl_options_client,
-                    cipher_list=self.config.ciphers_client,
-                    dhparams=self.config.certstore.dhparams,
-                    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:  # pragma: no cover
-            import traceback
-            self.log(
-                "Error in handle_sni:\r\n" +
-                traceback.format_exc(),
-                "error")
+    def log(self, msg, level):
+        msg = "{}: {}".format(repr(self.client_conn.address), msg)
+        self.channel.tell("log", Log(msg, level))
\ No newline at end of file