5 files changed, 103 insertions, 55 deletions
diff --git a/libmproxy/protocol/__init__.py b/libmproxy/protocol/__init__.py
index 0d624fd7..d8ebd4f0 100644
--- a/libmproxy/protocol/__init__.py
+++ b/libmproxy/protocol/__init__.py
@@ -28,12 +28,12 @@ as late as possible; this makes server replay without any outgoing connections p
 from __future__ import (absolute_import, print_function, division)
 from .base import Layer, ServerConnectionMixin, Kill
 from .http import Http1Layer, UpstreamConnectLayer, Http2Layer
-from .tls import TlsLayer, is_tls_record_magic
+from .tls import TlsLayer, is_tls_record_magic, TlsClientHello
 from .rawtcp import RawTCPLayer
 
 __all__ = [
     "Layer", "ServerConnectionMixin", "Kill",
     "Http1Layer", "UpstreamConnectLayer", "Http2Layer",
-    "TlsLayer", "is_tls_record_magic",
+    "TlsLayer", "is_tls_record_magic", "TlsClientHello"
     "RawTCPLayer"
 ]
diff --git a/libmproxy/protocol/base.py b/libmproxy/protocol/base.py
index af6b1c3b..d984cadb 100644
--- a/libmproxy/protocol/base.py
+++ b/libmproxy/protocol/base.py
@@ -111,7 +111,7 @@ class ServerConnectionMixin(object):
 
     def __init__(self, server_address=None):
         super(ServerConnectionMixin, self).__init__()
-        self.server_conn = ServerConnection(server_address)
+        self.server_conn = ServerConnection(server_address, (self.config.host, 0))
         self.__check_self_connect()
 
     def __check_self_connect(self):
@@ -157,10 +157,11 @@ class ServerConnectionMixin(object):
         """
         self.log("serverdisconnect", "debug", [repr(self.server_conn.address)])
         address = self.server_conn.address
+        source_address = self.server_conn.source_address
         self.server_conn.finish()
         self.server_conn.close()
         self.channel.tell("serverdisconnect", self.server_conn)
-        self.server_conn = ServerConnection(address)
+        self.server_conn = ServerConnection(address, source_address)
 
     def connect(self):
         """
diff --git a/libmproxy/protocol/http_replay.py b/libmproxy/protocol/http_replay.py
index b7faad07..63870dfb 100644
--- a/libmproxy/protocol/http_replay.py
+++ b/libmproxy/protocol/http_replay.py
@@ -46,7 +46,7 @@ class RequestReplayThread(threading.Thread):
                 # In all modes, we directly connect to the server displayed
                 if self.config.mode == "upstream":
                     server_address = self.config.upstream_server.address
-                    server = ServerConnection(server_address)
+                    server = ServerConnection(server_address, (self.config.host, 0))
                     server.connect()
                     if r.scheme == "https":
                         connect_request = make_connect_request((r.host, r.port))
@@ -68,7 +68,7 @@ class RequestReplayThread(threading.Thread):
                         r.form_out = "absolute"
                 else:
                     server_address = (r.host, r.port)
-                    server = ServerConnection(server_address)
+                    server = ServerConnection(server_address, (self.config.host, 0))
                     server.connect()
                     if r.scheme == "https":
                         server.establish_ssl(
diff --git a/libmproxy/protocol/rawtcp.py b/libmproxy/protocol/rawtcp.py
index 5f08fd17..ccd3c7ec 100644
--- a/libmproxy/protocol/rawtcp.py
+++ b/libmproxy/protocol/rawtcp.py
@@ -13,6 +13,15 @@ from ..exceptions import ProtocolException
 from .base import Layer
 
 
+class TcpMessage(object):
+    def __init__(self, client_conn, server_conn, sender, receiver, message):
+        self.client_conn = client_conn
+        self.server_conn = server_conn
+        self.sender = sender
+        self.receiver = receiver
+        self.message = message
+
+
 class RawTCPLayer(Layer):
     chunk_size = 4096
 
@@ -50,7 +59,13 @@ class RawTCPLayer(Layer):
                             return
                         continue
 
-                    dst.sendall(buf[:size])
+                    tcp_message = TcpMessage(
+                        self.client_conn, self.server_conn,
+                        self.client_conn if dst == server else self.server_conn,
+                        self.server_conn if dst == server else self.client_conn,
+                        buf[:size].tobytes())
+                    self.channel.ask("tcp_message", tcp_message)
+                    dst.sendall(tcp_message.message)
 
                     if self.logging:
                         # log messages are prepended with the client address,
@@ -59,7 +74,7 @@ class RawTCPLayer(Layer):
                             direction = "-> tcp -> {}".format(repr(self.server_conn.address))
                         else:
                             direction = "<- tcp <- {}".format(repr(self.server_conn.address))
-                        data = clean_bin(buf[:size].tobytes())
+                        data = clean_bin(tcp_message.message)
                         self.log(
                             "{}\r\n{}".format(direction, data),
                             "info"
diff --git a/libmproxy/protocol/tls.py b/libmproxy/protocol/tls.py
index ed747643..6d4cac85 100644
--- a/libmproxy/protocol/tls.py
+++ b/libmproxy/protocol/tls.py
@@ -221,6 +221,80 @@ def is_tls_record_magic(d):
         d[2] in ('\x00', '\x01', '\x02', '\x03')
     )
 
+def get_client_hello(client_conn):
+    """
+    Peek into the socket and read all records that contain the initial client hello message.
+
+    client_conn:
+        The :py:class:`client connection <libmproxy.models.ClientConnection>`.
+
+    Returns:
+        The raw handshake packet bytes, without TLS record header(s).
+    """
+    client_hello = ""
+    client_hello_size = 1
+    offset = 0
+    while len(client_hello) < client_hello_size:
+        record_header = client_conn.rfile.peek(offset + 5)[offset:]
+        if not is_tls_record_magic(record_header) or len(record_header) != 5:
+            raise TlsProtocolException('Expected TLS record, got "%s" instead.' % record_header)
+        record_size = struct.unpack("!H", record_header[3:])[0] + 5
+        record_body = client_conn.rfile.peek(offset + record_size)[offset + 5:]
+        if len(record_body) != record_size - 5:
+            raise TlsProtocolException("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
+    return client_hello
+
+class TlsClientHello(object):
+    def __init__(self, raw_client_hello):
+        self._client_hello = ClientHello.parse(raw_client_hello)
+
+    def raw(self):
+        return self._client_hello
+
+    @property
+    def client_cipher_suites(self):
+        return self._client_hello.cipher_suites.cipher_suites
+
+    @property
+    def client_sni(self):
+        for extension in self._client_hello.extensions:
+            if (extension.type == 0x00 and len(extension.server_names) == 1
+                    and extension.server_names[0].type == 0):
+                return extension.server_names[0].name
+
+    @property
+    def client_alpn_protocols(self):
+        for extension in self._client_hello.extensions:
+            if extension.type == 0x10:
+                return list(extension.alpn_protocols)
+
+    @classmethod
+    def from_client_conn(cls, client_conn):
+        """
+        Peek into the connection, read the initial client hello and parse it to obtain ALPN values.
+        client_conn:
+            The :py:class:`client connection <libmproxy.models.ClientConnection>`.
+        Returns:
+            :py:class:`client hello <libmproxy.protocol.tls.TlsClientHello>`.
+        """
+        try:
+            raw_client_hello = get_client_hello(client_conn)[4:]  # exclude handshake header.
+        except ProtocolException as e:
+            raise TlsProtocolException('Cannot read raw Client Hello: %s' % repr(e))
+
+        try:
+            return cls(raw_client_hello)
+        except ConstructError as e:
+            raise TlsProtocolException('Cannot parse Client Hello: %s, Raw Client Hello: %s' % \
+                            (repr(e), raw_client_hello.encode("hex")))
+
+    def __repr__(self):
+        return "TlsClientHello( sni: %s alpn_protocols: %s,  cipher_suites: %s)" % \
+                (self.client_sni, self.client_alpn_protocols, self.client_cipher_suites)
+
 
 class TlsLayer(Layer):
     def __init__(self, ctx, client_tls, server_tls):
@@ -281,60 +355,18 @@ class TlsLayer(Layer):
         else:
             return "TlsLayer(inactive)"
 
-    def _get_client_hello(self):
-        """
-        Peek into the socket and read all records that contain the initial client hello message.
-
-        Returns:
-            The raw handshake packet bytes, without TLS record header(s).
-        """
-        client_hello = ""
-        client_hello_size = 1
-        offset = 0
-        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 TlsProtocolException('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 TlsProtocolException("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
-        return client_hello
 
     def _parse_client_hello(self):
         """
         Peek into the connection, read the initial client hello and parse it to obtain ALPN values.
         """
         try:
-            raw_client_hello = self._get_client_hello()[4:]  # exclude handshake header.
-        except ProtocolException as e:
-            self.log("Cannot parse Client Hello: %s" % repr(e), "error")
-            return
-
-        try:
-            client_hello = ClientHello.parse(raw_client_hello)
-        except ConstructError as e:
+            parsed = TlsClientHello.from_client_conn(self.client_conn)
+            self.client_sni = parsed.client_sni
+            self.client_alpn_protocols = parsed.client_alpn_protocols
+            self.client_ciphers  = parsed.client_cipher_suites
+        except TlsProtocolException as e:
             self.log("Cannot parse Client Hello: %s" % repr(e), "error")
-            self.log("Raw Client Hello: %s" % raw_client_hello.encode("hex"), "debug")
-            return
-
-        self.client_ciphers = client_hello.cipher_suites.cipher_suites
-
-        for extension in client_hello.extensions:
-            if extension.type == 0x00:
-                if len(extension.server_names) != 1 or extension.server_names[0].type != 0:
-                    self.log("Unknown Server Name Indication: %s" % extension.server_names, "error")
-                self.client_sni = extension.server_names[0].name
-            elif extension.type == 0x10:
-                self.client_alpn_protocols = list(extension.alpn_protocols)
-
-        self.log(
-            "Parsed Client Hello: sni=%s, alpn=%s" % (self.client_sni, self.client_alpn_protocols),
-            "debug"
-        )
 
     def connect(self):
         if not self.server_conn: