split up proxy.py

9 files changed, 193 insertions, 178 deletions

diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index f8ad2444..c362465c 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -9,7 +9,6 @@ import flask
 import requests
 import tnetstring, filt, script
 from netlib import odict, wsgi
-from .proxy import ClientConnection, ServerConnection  # FIXME: remove circular dependency
 import controller, version, protocol
 import app
 from .protocol import KILL
diff --git a/libmproxy/protocol/__init__.py b/libmproxy/protocol/__init__.py
index 2c2e7285..392a8e3d 100644
--- a/libmproxy/protocol/__init__.py
+++ b/libmproxy/protocol/__init__.py
@@ -1,14 +1,7 @@
-from ..proxy import ServerConnection, AddressPriority
+from ..prxy.server import AddressPriority
 
 KILL = 0  # const for killed requests
 
-class ConnectionTypeChange(Exception):
-    """
-    Gets raised if the connetion type has been changed (e.g. after HTTP/1.1 101 Switching Protocols).
-    It's up to the raising ProtocolHandler to specify the new conntype before raising the exception.
-    """
-    pass
-
 
 class ProtocolHandler(object):
     def __init__(self, c):
diff --git a/libmproxy/protocol/http.py b/libmproxy/protocol/http.py
index 8a2583b1..3f9668ae 100644
--- a/libmproxy/protocol/http.py
+++ b/libmproxy/protocol/http.py
@@ -1,11 +1,13 @@
 import Cookie, urllib, urlparse, time, copy
 from email.utils import parsedate_tz, formatdate, mktime_tz
+from ..prxy.connection import ServerConnection
+from ..prxy.exception import ProxyError, ConnectionTypeChange
+from ..prxy.server import AddressPriority
 import netlib.utils
-from netlib import http, tcp, http_status, odict
+from netlib import http, tcp, http_status
 from netlib.odict import ODict, ODictCaseless
-from . import ProtocolHandler, ConnectionTypeChange, KILL, TemporaryServerChangeMixin
-from .. import encoding, utils, version, filt, controller, stateobject
-from ..proxy import ProxyError, AddressPriority, ServerConnection
+from . import ProtocolHandler, KILL, TemporaryServerChangeMixin
+from .. import encoding, utils, filt, controller, stateobject
 from .primitives import Flow, Error
 
 
diff --git a/libmproxy/protocol/primitives.py b/libmproxy/protocol/primitives.py
index 90191eeb..7cee074d 100644
--- a/libmproxy/protocol/primitives.py
+++ b/libmproxy/protocol/primitives.py
@@ -1,5 +1,5 @@
 from .. import stateobject, utils, version
-from ..proxy import ServerConnection, ClientConnection
+from ..prxy.connection import ClientConnection, ServerConnection
 import copy
 
 
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)
 
diff --git a/libmproxy/prxy/__init__.py b/libmproxy/prxy/__init__.py
new file mode 100644
index 00000000..e69de29b
--- /dev/null
+++ b/libmproxy/prxy/__init__.py
diff --git a/libmproxy/prxy/connection.py b/libmproxy/prxy/connection.py
new file mode 100644
index 00000000..b1040c1c
--- /dev/null
+++ b/libmproxy/prxy/connection.py
@@ -0,0 +1,138 @@
+import copy
+import os
+from .. import stateobject, utils
+from .exception import ProxyError
+from netlib import tcp, certutils
+
+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()
\ No newline at end of file
diff --git a/libmproxy/prxy/exception.py b/libmproxy/prxy/exception.py
new file mode 100644
index 00000000..c43a5d75
--- /dev/null
+++ b/libmproxy/prxy/exception.py
@@ -0,0 +1,14 @@
+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 ConnectionTypeChange(Exception):
+    """
+    Gets raised if the connection type has been changed (e.g. after HTTP/1.1 101 Switching Protocols).
+    It's up to the raising ProtocolHandler to specify the new conntype before raising the exception.
+    """
+    pass
\ No newline at end of file
diff --git a/libmproxy/prxy/server.py b/libmproxy/prxy/server.py
new file mode 100644
index 00000000..441b29b4
--- /dev/null
+++ b/libmproxy/prxy/server.py
@@ -0,0 +1,18 @@
+__author__ = 'user'
+
+
+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)"""
\ No newline at end of file