split up proxy.py

1 files changed, 15 insertions, 164 deletions

diff --git a/libmproxy/proxy.py b/libmproxy/proxy.py
index 6dd37752..ccb47c26 100644
--- a/libmproxy/proxy.py
+++ b/libmproxy/proxy.py
@@ -1,7 +1,18 @@
-import os, socket, time, threading, copy
+import os
+import socket
+import threading
+
 from OpenSSL import SSL
+
+from .prxy.connection import ClientConnection, ServerConnection
+from .prxy.exception import ProxyError, ConnectionTypeChange
+from .prxy.server import AddressPriority
 from netlib import tcp, http, certutils, http_auth
-import utils, version, platform, controller, stateobject
+import utils
+import version
+import platform
+import controller
+
 
 TRANSPARENT_SSL_PORTS = [443, 8443]
 CONF_BASENAME = "mitmproxy"
@@ -9,32 +20,6 @@ CONF_DIR = "~/.mitmproxy"
 CA_CERT_NAME = "mitmproxy-ca.pem"
 
 
-
-class AddressPriority(object):
-    """
-    Enum that signifies the priority of the given address when choosing the destination host.
-    Higher is better (None < i)
-    """
-    FORCE = 5
-    """forward mode"""
-    MANUALLY_CHANGED = 4
-    """user changed the target address in the ui"""
-    FROM_SETTINGS = 3
-    """reverse proxy mode"""
-    FROM_CONNECTION = 2
-    """derived from transparent resolver"""
-    FROM_PROTOCOL = 1
-    """derived from protocol (e.g. absolute-form http requests)"""
-
-
-class ProxyError(Exception):
-    def __init__(self, code, msg, headers=None):
-        self.code, self.msg, self.headers = code, msg, headers
-
-    def __str__(self):
-        return "ProxyError(%s, %s)" % (self.code, self.msg)
-
-
 class Log:
     def __init__(self, msg):
         self.msg = msg
@@ -58,140 +43,6 @@ class ProxyConfig:
         self.certstore = certutils.CertStore.from_store(self.confdir, CONF_BASENAME)
 
 
-
-class ClientConnection(tcp.BaseHandler, stateobject.SimpleStateObject):
-    def __init__(self, client_connection, address, server):
-        if client_connection:  # Eventually, this object is restored from state. We don't have a connection then.
-            tcp.BaseHandler.__init__(self, client_connection, address, server)
-        else:
-            self.connection = None
-            self.server = None
-            self.wfile = None
-            self.rfile = None
-            self.address = None
-            self.clientcert = None
-
-        self.timestamp_start = utils.timestamp()
-        self.timestamp_end = None
-        self.timestamp_ssl_setup = None
-
-    _stateobject_attributes = dict(
-        timestamp_start=float,
-        timestamp_end=float,
-        timestamp_ssl_setup=float
-    )
-
-    def _get_state(self):
-        d = super(ClientConnection, self)._get_state()
-        d.update(
-            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
-
-    def copy(self):
-        return copy.copy(self)
-
-    def send(self, message):
-        self.wfile.write(message)
-        self.wfile.flush()
-
-    @classmethod
-    def _from_state(cls, state):
-        f = cls(None, tuple(), None)
-        f._load_state(state)
-        return f
-
-    def convert_to_ssl(self, *args, **kwargs):
-        tcp.BaseHandler.convert_to_ssl(self, *args, **kwargs)
-        self.timestamp_ssl_setup = utils.timestamp()
-
-    def finish(self):
-        tcp.BaseHandler.finish(self)
-        self.timestamp_end = utils.timestamp()
-
-
-class ServerConnection(tcp.TCPClient, stateobject.SimpleStateObject):
-    def __init__(self, address, priority):
-        tcp.TCPClient.__init__(self, address)
-        self.priority = priority
-
-        self.peername = None
-        self.timestamp_start = None
-        self.timestamp_end = None
-        self.timestamp_tcp_setup = None
-        self.timestamp_ssl_setup = None
-
-    _stateobject_attributes = dict(
-        peername=tuple,
-        timestamp_start=float,
-        timestamp_end=float,
-        timestamp_tcp_setup=float,
-        timestamp_ssl_setup=float,
-        address=tcp.Address,
-        source_address=tcp.Address,
-        cert=certutils.SSLCert,
-        ssl_established=bool,
-        sni=str
-    )
-
-    def _get_state(self):
-        d = super(ServerConnection, self)._get_state()
-        d.update(
-            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,
-            cert=self.cert.to_pem() if self.cert else None
-        )
-        return d
-
-    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
-
-    @classmethod
-    def _from_state(cls, state):
-        f = cls(tuple(), None)
-        f._load_state(state)
-        return f
-
-    def copy(self):
-        return copy.copy(self)
-
-    def connect(self):
-        self.timestamp_start = utils.timestamp()
-        tcp.TCPClient.connect(self)
-        self.peername = self.connection.getpeername()
-        self.timestamp_tcp_setup = utils.timestamp()
-
-    def send(self, message):
-        self.wfile.write(message)
-        self.wfile.flush()
-
-    def establish_ssl(self, clientcerts, sni):
-        clientcert = None
-        if clientcerts:
-            path = os.path.join(clientcerts, self.address.host.encode("idna")) + ".pem"
-            if os.path.exists(path):
-                clientcert = path
-        try:
-            self.convert_to_ssl(cert=clientcert, sni=sni)
-            self.timestamp_ssl_setup = utils.timestamp()
-        except tcp.NetLibError, v:
-            raise ProxyError(400, str(v))
-
-    def finish(self):
-        tcp.TCPClient.finish(self)
-        self.timestamp_end = utils.timestamp()
-
 from . import protocol
 from .protocol.http import HTTPResponse
 
@@ -268,7 +119,7 @@ class ConnectionHandler:
                 while not self.close:
                     try:
                         protocol.handle_messages(self.conntype, self)
-                    except protocol.ConnectionTypeChange:
+                    except ConnectionTypeChange:
                         self.log("Connection Type Changed: %s" % self.conntype)
                         continue
 
@@ -323,7 +174,7 @@ class ConnectionHandler:
         """
         Sets a new server address with the given priority.
         Does not re-establish either connection or SSL handshake.
-        @type priority: AddressPriority
+        @type priority: libmproxy.prxy.server.AddressPriority
         """
         address = tcp.Address.wrap(address)