move files around

8 files changed, 215 insertions, 196 deletions

diff --git a/libmproxy/protocol2/__init__.py b/libmproxy/protocol2/__init__.py
new file mode 100644
index 00000000..9374a5bf
--- /dev/null
+++ b/libmproxy/protocol2/__init__.py
@@ -0,0 +1,157 @@
+"""
+mitmproxy protocol architecture
+
+In mitmproxy, protocols are implemented as a set of layers, which are composed on top each other.
+For example, the following scenarios depict possible scenarios (lowest layer first):
+
+Transparent HTTP proxy, no SSL:
+    TransparentModeLayer
+    HttpLayer
+
+Regular proxy, CONNECT request with WebSockets over SSL:
+    RegularModeLayer
+    HttpLayer
+    SslLayer
+    WebsocketLayer (or TcpLayer)
+
+Automated protocol detection by peeking into the buffer:
+    TransparentModeLayer
+    AutoLayer
+    SslLayer
+    AutoLayer
+    Http2Layer
+
+Communication between layers is done as follows:
+    - lower layers provide context information to higher layers
+    - higher layers can "yield" commands to lower layers,
+      which are propagated until they reach a suitable layer.
+
+Further goals:
+  - Connections should always be peekable to make automatic protocol detection work.
+  - Upstream connections should be established as late as possible;
+    inline scripts shall have a chance to handle everything locally.
+"""
+
+from __future__ import (absolute_import, print_function, division, unicode_literals)
+from netlib import tcp
+from ..proxy import ProxyError2, Log
+from ..proxy.connection import ServerConnection
+from .messages import *
+
+
+class RootContext(object):
+    """
+    The outmost context provided to the root layer.
+    As a consequence, every layer has .client_conn, .channel and .config.
+    """
+
+    def __init__(self, client_conn, config, channel):
+        self.client_conn = client_conn  # Client Connection
+        self.channel = channel  # provides .ask() method to communicate with FlowMaster
+        self.config = config  # Proxy Configuration
+
+    def __getattr__(self, name):
+        """
+        Accessing a nonexisting attribute does not throw an error but returns None instead.
+        """
+        return None
+
+
+class _LayerCodeCompletion(object):
+    """
+    Dummy class that provides type hinting in PyCharm, which simplifies development a lot.
+    """
+
+    def __init__(self):
+        if True:
+            return
+        self.config = None
+        """@type: libmproxy.proxy.config.ProxyConfig"""
+        self.client_conn = None
+        """@type: libmproxy.proxy.connection.ClientConnection"""
+        self.channel = None
+        """@type: libmproxy.controller.Channel"""
+
+
+class Layer(_LayerCodeCompletion):
+    def __init__(self, ctx):
+        """
+        Args:
+            ctx: The (read-only) higher layer.
+        """
+        super(Layer, self).__init__()
+        self.ctx = ctx
+
+    def __call__(self):
+        """
+        Logic of the layer.
+        Raises:
+            ProxyError2 in case of protocol exceptions.
+        """
+        raise NotImplementedError
+
+    def __getattr__(self, name):
+        """
+        Attributes not present on the current layer may exist on a higher layer.
+        """
+        return getattr(self.ctx, name)
+
+    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))
+
+
+class ServerConnectionMixin(object):
+    """
+    Mixin that provides a layer with the capabilities to manage a server connection.
+    """
+
+    def __init__(self):
+        self.server_address = None
+        self.server_conn = None
+
+    def _handle_server_message(self, message):
+        if message == Reconnect:
+            self._disconnect()
+            self._connect()
+            return True
+        elif message == Connect:
+            self._connect()
+            return True
+        elif message == ChangeServer:
+            raise NotImplementedError
+        return False
+
+    def _disconnect(self):
+        """
+        Deletes (and closes) an existing server connection.
+        """
+        self.log("serverdisconnect", "debug", [repr(self.server_address)])
+        self.server_conn.finish()
+        self.server_conn.close()
+        # self.channel.tell("serverdisconnect", self)
+        self.server_conn = None
+
+    def _connect(self):
+        self.log("serverconnect", "debug", [repr(self.server_address)])
+        self.server_conn = ServerConnection(self.server_address)
+        try:
+            self.server_conn.connect()
+        except tcp.NetLibError as e:
+            raise ProxyError2("Server connection to '%s' failed: %s" % (self.server_address, e), e)
+
+    def _set_address(self, address):
+        a = tcp.Address.wrap(address)
+        self.log("Set new server address: " + repr(a), "debug")
+        self.server_address = address
+
+
+from .socks import Socks5IncomingLayer
+from .rawtcp import TcpLayer
\ No newline at end of file
diff --git a/libmproxy/proxy/message.py b/libmproxy/protocol2/messages.py
index 7eb59344..52bb5a44 100644
--- a/libmproxy/proxy/message.py
+++ b/libmproxy/protocol2/messages.py
@@ -1,6 +1,7 @@
 """
 This module contains all valid messages layers can send to the underlying layers.
 """
+from __future__ import (absolute_import, print_function, division, unicode_literals)
 
 
 class _Message(object):
diff --git a/libmproxy/protocol2/rawtcp.py b/libmproxy/protocol2/rawtcp.py
new file mode 100644
index 00000000..db9a48fa
--- /dev/null
+++ b/libmproxy/protocol2/rawtcp.py
@@ -0,0 +1,13 @@
+from . import Layer, Connect
+from ..protocol.tcp import TCPHandler
+
+
+class TcpLayer(Layer):
+    def __call__(self):
+        yield Connect()
+        tcp_handler = TCPHandler(self)
+        tcp_handler.handle_messages()
+
+    def establish_server_connection(self):
+        pass
+        # FIXME: Remove method, currently just here to mock TCPHandler's call to it.
diff --git a/libmproxy/protocol2/socks.py b/libmproxy/protocol2/socks.py
new file mode 100644
index 00000000..90771015
--- /dev/null
+++ b/libmproxy/protocol2/socks.py
@@ -0,0 +1,26 @@
+from __future__ import (absolute_import, print_function, division, unicode_literals)
+
+from ..proxy import ProxyError, Socks5ProxyMode, ProxyError2
+from . import Layer, ServerConnectionMixin
+from .rawtcp import TcpLayer
+from .ssl import SslLayer
+
+
+class Socks5IncomingLayer(Layer, ServerConnectionMixin):
+    def __call__(self):
+        try:
+            s5mode = Socks5ProxyMode(self.config.ssl_ports)
+            address = s5mode.get_upstream_server(self.client_conn)[2:]
+        except ProxyError as e:
+            # TODO: Unmonkeypatch
+            raise ProxyError2(str(e), e)
+
+        self._set_address(address)
+
+        if address[1] == 443:
+            layer = SslLayer(self, True, True)
+        else:
+            layer = TcpLayer(self)
+        for message in layer():
+            if not self._handle_server_message(message):
+                yield message
diff --git a/libmproxy/proxy/layer.py b/libmproxy/protocol2/ssl.py
index e4878bdf..6b44bf42 100644
--- a/libmproxy/proxy/layer.py
+++ b/libmproxy/protocol2/ssl.py
@@ -2,196 +2,12 @@ from __future__ import (absolute_import, print_function, division, unicode_liter
 import Queue
 import threading
 import traceback
-from libmproxy.protocol.tcp import TCPHandler
-from libmproxy.proxy.connection import ServerConnection
 from netlib import tcp
-from .primitives import Socks5ProxyMode, ProxyError, Log
-from .message import Connect, Reconnect, ChangeServer
 
-"""
-mitmproxy protocol architecture
-
-In mitmproxy, protocols are implemented as a set of layers, which are composed on top each other.
-For example, the following scenarios depict possible scenarios (lowest layer first):
-
-Transparent HTTP proxy, no SSL:
-    TransparentModeLayer
-    HttpLayer
-
-Regular proxy, CONNECT request with WebSockets over SSL:
-    RegularModeLayer
-    HttpLayer
-    SslLayer
-    WebsocketLayer (or TcpLayer)
-
-Automated protocol detection by peeking into the buffer:
-    TransparentModeLayer
-    AutoLayer
-    SslLayer
-    AutoLayer
-    Http2Layer
-
-Communication between layers is done as follows:
-    - lower layers provide context information to higher layers
-    - higher layers can "yield" commands to lower layers,
-      which are propagated until they reach a suitable layer.
-
-Further goals:
-  - Connections should always be peekable to make automatic protocol detection work.
-  - Upstream connections should be established as late as possible;
-    inline scripts shall have a chance to handle everything locally.
-"""
-
-
-class ProxyError2(Exception):
-    def __init__(self, message, cause=None):
-        super(ProxyError2, self).__init__(message)
-        self.cause = cause
-
-
-class RootContext(object):
-    """
-    The outmost context provided to the root layer.
-    As a consequence, every layer has .client_conn, .channel and .config.
-    """
-
-    def __init__(self, client_conn, config, channel):
-        self.client_conn = client_conn  # Client Connection
-        self.channel = channel  # provides .ask() method to communicate with FlowMaster
-        self.config = config  # Proxy Configuration
-
-    def __getattr__(self, name):
-        """
-        Accessing a nonexisting attribute does not throw an error but returns None instead.
-        """
-        return None
-
-
-class _LayerCodeCompletion(object):
-    """
-    Dummy class that provides type hinting in PyCharm, which simplifies development a lot.
-    """
-
-    def __init__(self):
-        if True:
-            return
-        self.config = None
-        """@type: libmproxy.proxy.config.ProxyConfig"""
-        self.client_conn = None
-        """@type: libmproxy.proxy.connection.ClientConnection"""
-        self.channel = None
-        """@type: libmproxy.controller.Channel"""
-
-
-class Layer(_LayerCodeCompletion):
-    def __init__(self, ctx):
-        """
-        Args:
-            ctx: The (read-only) higher layer.
-        """
-        super(Layer, self).__init__()
-        self.ctx = ctx
-
-    def __call__(self):
-        """
-        Logic of the layer.
-        Raises:
-            ProxyError2 in case of protocol exceptions.
-        """
-        raise NotImplementedError
-
-    def __getattr__(self, name):
-        """
-        Attributes not present on the current layer may exist on a higher layer.
-        """
-        return getattr(self.ctx, name)
-
-    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))
-
-
-class _ServerConnectionMixin(object):
-    """
-    Mixin that provides a layer with the capabilities to manage a server connection.
-    """
-
-    def __init__(self):
-        self.server_address = None
-        self.server_conn = None
-
-    def _handle_server_message(self, message):
-        if message == Reconnect:
-            self._disconnect()
-            self._connect()
-            return True
-        elif message == Connect:
-            self._connect()
-            return True
-        elif message == ChangeServer:
-            raise NotImplementedError
-        return False
-
-    def _disconnect(self):
-        """
-        Deletes (and closes) an existing server connection.
-        """
-        self.log("serverdisconnect", "debug", [repr(self.server_address)])
-        self.server_conn.finish()
-        self.server_conn.close()
-        # self.channel.tell("serverdisconnect", self)
-        self.server_conn = None
-
-    def _connect(self):
-        self.log("serverconnect", "debug", [repr(self.server_address)])
-        self.server_conn = ServerConnection(self.server_address)
-        try:
-            self.server_conn.connect()
-        except tcp.NetLibError as e:
-            raise ProxyError2("Server connection to '%s' failed: %s" % (self.server_address, e), e)
-
-    def _set_address(self, address):
-        a = tcp.Address.wrap(address)
-        self.log("Set new server address: " + repr(a), "debug")
-        self.server_address = address
-
-
-class Socks5IncomingLayer(Layer, _ServerConnectionMixin):
-    def __call__(self):
-        try:
-            s5mode = Socks5ProxyMode(self.config.ssl_ports)
-            address = s5mode.get_upstream_server(self.client_conn)[2:]
-        except ProxyError as e:
-            # TODO: Unmonkeypatch
-            raise ProxyError2(str(e), e)
-
-        self._set_address(address)
-
-        if address[1] == 443:
-            layer = SslLayer(self, True, True)
-        else:
-            layer = TcpLayer(self)
-        for message in layer():
-            if not self._handle_server_message(message):
-                yield message
-
-
-class TcpLayer(Layer):
-    def __call__(self):
-        yield Connect()
-        tcp_handler = TCPHandler(self)
-        tcp_handler.handle_messages()
-
-    def establish_server_connection(self):
-        pass
-        # FIXME: Remove method, currently just here to mock TCPHandler's call to it.
+from ..proxy import ProxyError2
+from . import Layer
+from .messages import Connect, Reconnect, ChangeServer
+from .rawtcp import TcpLayer
 
 
 class ReconnectRequest(object):
diff --git a/libmproxy/proxy/__init__.py b/libmproxy/proxy/__init__.py
index f33d323b..7d664707 100644
--- a/libmproxy/proxy/__init__.py
+++ b/libmproxy/proxy/__init__.py
@@ -1,2 +1,2 @@
 from .primitives import *
-from .config import ProxyConfig, process_proxy_options
+from .config import ProxyConfig, process_proxy_options
\ No newline at end of file
diff --git a/libmproxy/proxy/primitives.py b/libmproxy/proxy/primitives.py
index a9f31181..fd4eb882 100644
--- a/libmproxy/proxy/primitives.py
+++ b/libmproxy/proxy/primitives.py
@@ -2,6 +2,12 @@ from __future__ import absolute_import
 from netlib import socks, tcp
 
 
+class ProxyError2(Exception):
+    def __init__(self, message, cause=None):
+        super(ProxyError2, self).__init__(message)
+        self.cause = cause
+
+
 class ProxyError(Exception):
     def __init__(self, code, message, headers=None):
         super(ProxyError, self).__init__(message)
diff --git a/libmproxy/proxy/server.py b/libmproxy/proxy/server.py
index a45276d4..c8990a9a 100644
--- a/libmproxy/proxy/server.py
+++ b/libmproxy/proxy/server.py
@@ -3,12 +3,12 @@ from __future__ import absolute_import, print_function
 import traceback
 import sys
 import socket
-
 from netlib import tcp
-from . import layer
-from .primitives import ProxyServerError, Log, ProxyError
-from .connection import ClientConnection, ServerConnection
+
 from ..protocol.handle import protocol_handler
+from .. import protocol2
+from .primitives import ProxyServerError, Log, ProxyError, ProxyError2
+from .connection import ClientConnection, ServerConnection
 
 
 class DummyServer:
@@ -74,17 +74,17 @@ class ConnectionHandler2:
     def handle(self):
         self.log("clientconnect", "info")
 
-        root_context = layer.RootContext(
+        root_context = protocol2.RootContext(
             self.client_conn,
             self.config,
             self.channel
         )
-        root_layer = layer.Socks5IncomingLayer(root_context)
+        root_layer = protocol2.Socks5IncomingLayer(root_context)
 
         try:
             for message in root_layer():
                 print("Root layer receveived: %s" % message)
-        except layer.ProxyError2 as e:
+        except ProxyError2 as e:
             self.log(e, "info")
         except Exception:
             self.log(traceback.format_exc(), "error")