add auto layer, multiple other fixes

6 files changed, 183 insertions, 165 deletions

diff --git a/libmproxy/protocol2/__init__.py b/libmproxy/protocol2/__init__.py
index 9374a5bf..6f4bfe44 100644
--- a/libmproxy/protocol2/__init__.py
+++ b/libmproxy/protocol2/__init__.py
@@ -1,157 +1,6 @@
-"""
-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 .layer import RootContext
 from .socks import Socks5IncomingLayer
-from .rawtcp import TcpLayer
\ No newline at end of file
+from .rawtcp import TcpLayer
+from .auto import AutoLayer
+__all__ = ["Socks5IncomingLayer", "TcpLayer", "AutoLayer", "RootContext"]
\ No newline at end of file
diff --git a/libmproxy/protocol2/auto.py b/libmproxy/protocol2/auto.py
new file mode 100644
index 00000000..1c4293ac
--- /dev/null
+++ b/libmproxy/protocol2/auto.py
@@ -0,0 +1,16 @@
+from __future__ import (absolute_import, print_function, division, unicode_literals)
+from .layer import Layer
+
+
+class AutoLayer(Layer):
+    def __call__(self):
+        d = self.client_conn.rfile.peek(1)
+        if d[0] == "\x16":
+            layer = SslLayer(self, True, True)
+        else:
+            layer = TcpLayer(self)
+        for m in layer():
+            yield m
+
+from .rawtcp import TcpLayer
+from .ssl import SslLayer
diff --git a/libmproxy/protocol2/layer.py b/libmproxy/protocol2/layer.py
new file mode 100644
index 00000000..14263f64
--- /dev/null
+++ b/libmproxy/protocol2/layer.py
@@ -0,0 +1,151 @@
+"""
+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 Connect, Reconnect, ChangeServer
+
+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
\ No newline at end of file
diff --git a/libmproxy/protocol2/rawtcp.py b/libmproxy/protocol2/rawtcp.py
index db9a48fa..b40c569f 100644
--- a/libmproxy/protocol2/rawtcp.py
+++ b/libmproxy/protocol2/rawtcp.py
@@ -1,5 +1,7 @@
-from . import Layer, Connect
+from __future__ import (absolute_import, print_function, division, unicode_literals)
 from ..protocol.tcp import TCPHandler
+from .layer import Layer
+from .messages import Connect
 
 
 class TcpLayer(Layer):
diff --git a/libmproxy/protocol2/socks.py b/libmproxy/protocol2/socks.py
index 90771015..9ca30bb4 100644
--- a/libmproxy/protocol2/socks.py
+++ b/libmproxy/protocol2/socks.py
@@ -1,9 +1,7 @@
 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
+from .layer import Layer, ServerConnectionMixin
 
 
 class Socks5IncomingLayer(Layer, ServerConnectionMixin):
@@ -18,9 +16,11 @@ class Socks5IncomingLayer(Layer, ServerConnectionMixin):
         self._set_address(address)
 
         if address[1] == 443:
-            layer = SslLayer(self, True, True)
+            layer = AutoLayer(self)
         else:
-            layer = TcpLayer(self)
+            layer = AutoLayer(self)
         for message in layer():
             if not self._handle_server_message(message):
                 yield message
+
+from .auto import AutoLayer
diff --git a/libmproxy/protocol2/ssl.py b/libmproxy/protocol2/ssl.py
index 6b44bf42..e8ff16cf 100644
--- a/libmproxy/protocol2/ssl.py
+++ b/libmproxy/protocol2/ssl.py
@@ -5,9 +5,9 @@ import traceback
 from netlib import tcp
 
 from ..proxy import ProxyError2
-from . import Layer
+from .layer import Layer
 from .messages import Connect, Reconnect, ChangeServer
-from .rawtcp import TcpLayer
+from .auto import AutoLayer
 
 
 class ReconnectRequest(object):
@@ -61,7 +61,7 @@ class SslLayer(Layer):
         elif self.client_ssl:
             self._establish_ssl_with_client()
 
-        layer = TcpLayer(self)
+        layer = AutoLayer(self)
         for message in layer():
             if message != Connect or not self._connected:
                 yield message
@@ -225,4 +225,4 @@ class SslLayer(Layer):
         if self._sni_from_server_change:
             sans.append(self._sni_from_server_change)
 
-        return self.config.certstore.get_cert(host, sans)
+        return self.config.certstore.get_cert(host, sans)
\ No newline at end of file