mitmproxy/libmproxy/proxy/connection.py
2014-06-25 23:03:15 +02:00

143 lines
4.6 KiB
Python

from __future__ import absolute_import
import copy
import os
from netlib import tcp, certutils
from .. import stateobject, utils
from .primitives import ProxyError
class ClientConnection(tcp.BaseHandler, stateobject.SimpleStateObject):
def __init__(self, client_connection, address, server):
if client_connection: # Eventually, this object is restored from state. We don't have a connection then.
tcp.BaseHandler.__init__(self, client_connection, address, server)
else:
self.connection = None
self.server = None
self.wfile = None
self.rfile = None
self.address = None
self.clientcert = None
self.timestamp_start = utils.timestamp()
self.timestamp_end = None
self.timestamp_ssl_setup = None
_stateobject_attributes = dict(
timestamp_start=float,
timestamp_end=float,
timestamp_ssl_setup=float
)
def _get_state(self):
d = super(ClientConnection, self)._get_state()
d.update(
address={"address": self.address(), "use_ipv6": self.address.use_ipv6},
clientcert=self.cert.to_pem() if self.clientcert else None
)
return d
def _load_state(self, state):
super(ClientConnection, self)._load_state(state)
self.address = tcp.Address(**state["address"]) if state["address"] else None
self.clientcert = certutils.SSLCert.from_pem(state["clientcert"]) if state["clientcert"] else None
def copy(self):
return copy.copy(self)
def send(self, message):
self.wfile.write(message)
self.wfile.flush()
@classmethod
def _from_state(cls, state):
f = cls(None, tuple(), None)
f._load_state(state)
return f
def convert_to_ssl(self, *args, **kwargs):
tcp.BaseHandler.convert_to_ssl(self, *args, **kwargs)
self.timestamp_ssl_setup = utils.timestamp()
def finish(self):
tcp.BaseHandler.finish(self)
self.timestamp_end = utils.timestamp()
class ServerConnection(tcp.TCPClient, stateobject.SimpleStateObject):
def __init__(self, address, priority):
tcp.TCPClient.__init__(self, address)
self.priority = priority
self.peername = None
self.sockname = None
self.timestamp_start = None
self.timestamp_end = None
self.timestamp_tcp_setup = None
self.timestamp_ssl_setup = None
_stateobject_attributes = dict(
peername=tuple,
sockname=tuple,
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
)
def _get_state(self):
d = super(ServerConnection, self)._get_state()
d.update(
address={"address": self.address(), "use_ipv6": self.address.use_ipv6},
source_address= {"address": self.source_address(),
"use_ipv6": self.source_address.use_ipv6} if self.source_address else None,
cert=self.cert.to_pem() if self.cert else None
)
return d
def _load_state(self, state):
super(ServerConnection, self)._load_state(state)
self.address = tcp.Address(**state["address"]) if state["address"] else None
self.source_address = tcp.Address(**state["source_address"]) if state["source_address"] else None
self.cert = certutils.SSLCert.from_pem(state["cert"]) if state["cert"] else None
@classmethod
def _from_state(cls, state):
f = cls(tuple(), None)
f._load_state(state)
return f
def copy(self):
return copy.copy(self)
def connect(self):
self.timestamp_start = utils.timestamp()
tcp.TCPClient.connect(self)
self.peername = self.connection.getpeername()
self.sockname = self.connection.getsockname()
self.timestamp_tcp_setup = utils.timestamp()
def send(self, message):
self.wfile.write(message)
self.wfile.flush()
def establish_ssl(self, clientcerts, sni):
clientcert = None
if clientcerts:
path = os.path.join(clientcerts, self.address.host.encode("idna")) + ".pem"
if os.path.exists(path):
clientcert = path
try:
self.convert_to_ssl(cert=clientcert, sni=sni)
self.timestamp_ssl_setup = utils.timestamp()
except tcp.NetLibError, v:
raise ProxyError(400, str(v))
def finish(self):
tcp.TCPClient.finish(self)
self.timestamp_end = utils.timestamp()