1 files changed, 161 insertions, 0 deletions
diff --git a/mitmproxy/libmproxy/models/connections.py b/mitmproxy/libmproxy/models/connections.py
new file mode 100644
index 00000000..d5920256
--- /dev/null
+++ b/mitmproxy/libmproxy/models/connections.py
@@ -0,0 +1,161 @@
+from __future__ import (absolute_import, print_function, division)
+
+import copy
+import os
+
+from netlib import tcp, certutils
+from .. import stateobject, utils
+
+
+class ClientConnection(tcp.BaseHandler, stateobject.StateObject):
+
+    def __init__(self, client_connection, address, server):
+        # Eventually, this object is restored from state. We don't have a
+        # connection then.
+        if client_connection:
+            super(ClientConnection, self).__init__(client_connection, address, server)
+        else:
+            self.connection = None
+            self.server = None
+            self.wfile = None
+            self.rfile = None
+            self.address = None
+            self.clientcert = None
+            self.ssl_established = None
+
+        self.timestamp_start = utils.timestamp()
+        self.timestamp_end = None
+        self.timestamp_ssl_setup = None
+        self.protocol = None
+
+    def __nonzero__(self):
+        return bool(self.connection) and not self.finished
+
+    def __repr__(self):
+        return "<ClientConnection: {ssl}{address}>".format(
+            ssl="[ssl] " if self.ssl_established else "",
+            address=repr(self.address)
+        )
+
+    @property
+    def tls_established(self):
+        return self.ssl_established
+
+    _stateobject_attributes = dict(
+        address=tcp.Address,
+        clientcert=certutils.SSLCert,
+        ssl_established=bool,
+        timestamp_start=float,
+        timestamp_end=float,
+        timestamp_ssl_setup=float
+    )
+
+    def copy(self):
+        return copy.copy(self)
+
+    def send(self, message):
+        if isinstance(message, list):
+            message = b''.join(message)
+        self.wfile.write(message)
+        self.wfile.flush()
+
+    @classmethod
+    def from_state(cls, state):
+        f = cls(None, tuple(), None)
+        f.set_state(state)
+        return f
+
+    def convert_to_ssl(self, *args, **kwargs):
+        super(ClientConnection, self).convert_to_ssl(*args, **kwargs)
+        self.timestamp_ssl_setup = utils.timestamp()
+
+    def finish(self):
+        super(ClientConnection, self).finish()
+        self.timestamp_end = utils.timestamp()
+
+
+class ServerConnection(tcp.TCPClient, stateobject.StateObject):
+
+    def __init__(self, address, source_address=None):
+        tcp.TCPClient.__init__(self, address, source_address)
+
+        self.via = None
+        self.timestamp_start = None
+        self.timestamp_end = None
+        self.timestamp_tcp_setup = None
+        self.timestamp_ssl_setup = None
+        self.protocol = None
+
+    def __nonzero__(self):
+        return bool(self.connection) and not self.finished
+
+    def __repr__(self):
+        if self.ssl_established and self.sni:
+            ssl = "[ssl: {0}] ".format(self.sni)
+        elif self.ssl_established:
+            ssl = "[ssl] "
+        else:
+            ssl = ""
+        return "<ServerConnection: {ssl}{address}>".format(
+            ssl=ssl,
+            address=repr(self.address)
+        )
+
+    @property
+    def tls_established(self):
+        return self.ssl_established
+
+    _stateobject_attributes = dict(
+        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
+    )
+
+    @classmethod
+    def from_state(cls, state):
+        f = cls(tuple())
+        f.set_state(state)
+        return f
+
+    def copy(self):
+        return copy.copy(self)
+
+    def connect(self):
+        self.timestamp_start = utils.timestamp()
+        tcp.TCPClient.connect(self)
+        self.timestamp_tcp_setup = utils.timestamp()
+
+    def send(self, message):
+        if isinstance(message, list):
+            message = b''.join(message)
+        self.wfile.write(message)
+        self.wfile.flush()
+
+    def establish_ssl(self, clientcerts, sni, **kwargs):
+        clientcert = None
+        if clientcerts:
+            if os.path.isfile(clientcerts):
+                clientcert = clientcerts
+            else:
+                path = os.path.join(
+                    clientcerts,
+                    self.address.host.encode("idna")) + ".pem"
+                if os.path.exists(path):
+                    clientcert = path
+
+        self.convert_to_ssl(cert=clientcert, sni=sni, **kwargs)
+        self.sni = sni
+        self.timestamp_ssl_setup = utils.timestamp()
+
+    def finish(self):
+        tcp.TCPClient.finish(self)
+        self.timestamp_end = utils.timestamp()
+
+
+ServerConnection._stateobject_attributes["via"] = ServerConnection