expose `next_layer` to inline scripts

6 files changed, 33 insertions, 23 deletions

diff --git a/libmproxy/exceptions.py b/libmproxy/exceptions.py
index 59436df7..6b997041 100644
--- a/libmproxy/exceptions.py
+++ b/libmproxy/exceptions.py
@@ -29,6 +29,10 @@ class ProtocolException(ProxyException):
     pass
 
 
+class TlsException(ProtocolException):
+    pass
+
+
 class Socks5Exception(ProtocolException):
     pass
 
diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index 547d0f60..d037d36e 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -945,21 +945,25 @@ class FlowMaster(controller.Master):
         self.add_event(l.msg, l.level)
         l.reply()
 
-    def handle_clientconnect(self, cc):
-        self.run_script_hook("clientconnect", cc)
-        cc.reply()
+    def handle_clientconnect(self, root_layer):
+        self.run_script_hook("clientconnect", root_layer)
+        root_layer.reply()
 
-    def handle_clientdisconnect(self, r):
-        self.run_script_hook("clientdisconnect", r)
-        r.reply()
+    def handle_clientdisconnect(self, root_layer):
+        self.run_script_hook("clientdisconnect", root_layer)
+        root_layer.reply()
 
-    def handle_serverconnect(self, sc):
-        self.run_script_hook("serverconnect", sc)
-        sc.reply()
+    def handle_serverconnect(self, server_conn):
+        self.run_script_hook("serverconnect", server_conn)
+        server_conn.reply()
 
-    def handle_serverdisconnect(self, sc):
-        self.run_script_hook("serverdisconnect", sc)
-        sc.reply()
+    def handle_serverdisconnect(self, server_conn):
+        self.run_script_hook("serverdisconnect", server_conn)
+        server_conn.reply()
+
+    def handle_next_layer(self, top_layer):
+        self.run_script_hook("next_layer", top_layer)
+        top_layer.reply()
 
     def handle_error(self, f):
         self.state.update_flow(f)
diff --git a/libmproxy/protocol/http.py b/libmproxy/protocol/http.py
index fbf4ac9b..93972111 100644
--- a/libmproxy/protocol/http.py
+++ b/libmproxy/protocol/http.py
@@ -3,7 +3,6 @@ from __future__ import (absolute_import, print_function, division)
 from netlib import tcp
 from netlib.http import http1, HttpErrorConnClosed, HttpError, Headers
 from netlib.http.semantics import CONTENT_MISSING
-from netlib import odict
 from netlib.tcp import NetLibError, Address
 from netlib.http.http1 import HTTP1Protocol
 from netlib.http.http2 import HTTP2Protocol
diff --git a/libmproxy/protocol/tls.py b/libmproxy/protocol/tls.py
index 88a8398f..6e8535ae 100644
--- a/libmproxy/protocol/tls.py
+++ b/libmproxy/protocol/tls.py
@@ -6,8 +6,8 @@ from construct import ConstructError
 
 from netlib.tcp import NetLibError, NetLibInvalidCertificateError
 from netlib.http.http1 import HTTP1Protocol
-from ..contrib.tls._constructs import ClientHello, CipherSuites
-from ..exceptions import ProtocolException
+from ..contrib.tls._constructs import ClientHello
+from ..exceptions import ProtocolException, TlsException
 from .base import Layer
 
 
@@ -201,6 +201,7 @@ CIPHER_ID_NAME_MAP = {
     0x080080: 'RC4-64-MD5',
 }
 
+
 def is_tls_record_magic(d):
     """
     Returns:
@@ -290,11 +291,11 @@ class TlsLayer(Layer):
         while len(client_hello) < client_hello_size:
             record_header = self.client_conn.rfile.peek(offset + 5)[offset:]
             if not is_tls_record_magic(record_header) or len(record_header) != 5:
-                raise ProtocolException('Expected TLS record, got "%s" instead.' % record_header)
+                raise TlsException('Expected TLS record, got "%s" instead.' % record_header)
             record_size = struct.unpack("!H", record_header[3:])[0] + 5
             record_body = self.client_conn.rfile.peek(offset + record_size)[offset + 5:]
             if len(record_body) != record_size - 5:
-                raise ProtocolException("Unexpected EOF in TLS handshake: %s" % record_body)
+                raise TlsException("Unexpected EOF in TLS handshake: %s" % record_body)
             client_hello += record_body
             offset += record_size
             client_hello_size = struct.unpack("!I", '\x00' + client_hello[1:4])[0] + 4
@@ -405,7 +406,7 @@ class TlsLayer(Layer):
                 alpn_select_callback=self.__alpn_select_callback,
             )
         except NetLibError as e:
-            raise ProtocolException("Cannot establish TLS with client: %s" % repr(e), e)
+            raise TlsException("Cannot establish TLS with client: %s" % repr(e), e)
 
     def _establish_tls_with_server(self):
         self.log("Establish TLS with server", "debug")
@@ -452,13 +453,13 @@ class TlsLayer(Layer):
                 (tls_cert_err['depth'], tls_cert_err['errno']),
                 "error")
             self.log("Aborting connection attempt", "error")
-            raise ProtocolException("Cannot establish TLS with {address} (sni: {sni}): {e}".format(
+            raise TlsException("Cannot establish TLS with {address} (sni: {sni}): {e}".format(
                 address=repr(self.server_conn.address),
                 sni=self.sni_for_server_connection,
                 e=repr(e),
             ), e)
         except NetLibError as e:
-            raise ProtocolException("Cannot establish TLS with {address} (sni: {sni}): {e}".format(
+            raise TlsException("Cannot establish TLS with {address} (sni: {sni}): {e}".format(
                 address=repr(self.server_conn.address),
                 sni=self.sni_for_server_connection,
                 e=repr(e),
@@ -487,5 +488,4 @@ class TlsLayer(Layer):
         if self._sni_from_server_change:
             sans.add(self._sni_from_server_change)
 
-        sans.discard(host)
         return self.config.certstore.get_cert(host, list(sans))
diff --git a/libmproxy/proxy/root_context.py b/libmproxy/proxy/root_context.py
index 87a540c0..dccdf023 100644
--- a/libmproxy/proxy/root_context.py
+++ b/libmproxy/proxy/root_context.py
@@ -40,7 +40,10 @@ class RootContext(object):
         Returns:
             The next layer
         """
+        layer = self._next_layer(top_layer)
+        return self.channel.ask("next_layer", layer)
 
+    def _next_layer(self, top_layer):
         # 1. Check for --ignore.
         if self.config.check_ignore(top_layer.server_conn.address):
             return RawTCPLayer(top_layer, logging=False)
@@ -119,4 +122,4 @@ class RootContext(object):
 class Log(object):
     def __init__(self, msg, level="info"):
         self.msg = msg
-        self.level = level
\ No newline at end of file
+        self.level = level
diff --git a/libmproxy/proxy/server.py b/libmproxy/proxy/server.py
index 5d067b45..c12bbbfa 100644
--- a/libmproxy/proxy/server.py
+++ b/libmproxy/proxy/server.py
@@ -137,4 +137,4 @@ class ConnectionHandler(object):
 
     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
+        self.channel.tell("log", Log(msg, level))