split up protocol\__init__.py

This commit is contained in:
Maximilian Hils 2014-03-10 21:57:50 +01:00
parent b59013f6e3
commit e6349b540f
8 changed files with 113 additions and 116 deletions

View File

@ -767,7 +767,7 @@ class FlowReader:
v = ".".join(str(i) for i in data["version"]) v = ".".join(str(i) for i in data["version"])
raise FlowReadError("Incompatible serialized data version: %s"%v) raise FlowReadError("Incompatible serialized data version: %s"%v)
off = self.fo.tell() off = self.fo.tell()
yield protocol.protocols[data["conntype"]]["flow"]._from_state(data) yield protocol.handle.protocols[data["conntype"]]["flow"]._from_state(data)
except ValueError, v: except ValueError, v:
# Error is due to EOF # Error is due to EOF
if self.fo.tell() == off and self.fo.read() == '': if self.fo.tell() == off and self.fo.read() == '':

View File

@ -1,95 +1 @@
from ..proxy.primitives import AddressPriority from .primitives import *
KILL = 0 # const for killed requests
class ProtocolHandler(object):
def __init__(self, c):
self.c = c
"""@type: libmproxy.proxy.ConnectionHandler"""
def handle_messages(self):
"""
This method gets called if a client connection has been made. Depending on the proxy settings,
a server connection might already exist as well.
"""
raise NotImplementedError # pragma: nocover
def handle_error(self, error):
"""
This method gets called should there be an uncaught exception during the connection.
This might happen outside of handle_messages, e.g. if the initial SSL handshake fails in transparent mode.
"""
raise error # pragma: nocover
class TemporaryServerChangeMixin(object):
"""
This mixin allows safe modification of the target server,
without any need to expose the ConnectionHandler to the Flow.
"""
def change_server(self, address, ssl):
if address == self.c.server_conn.address():
return
priority = AddressPriority.MANUALLY_CHANGED
if self.c.server_conn.priority > priority:
self.log("Attempt to change server address, "
"but priority is too low (is: %s, got: %s)" % (self.server_conn.priority, priority))
return
self.log("Temporarily change server connection: %s:%s -> %s:%s" % (
self.c.server_conn.address.host,
self.c.server_conn.address.port,
address.host,
address.port
))
if not hasattr(self, "_backup_server_conn"):
self._backup_server_conn = self.c.server_conn
self.c.server_conn = None
else: # This is at least the second temporary change. We can kill the current connection.
self.c.del_server_connection()
self.c.set_server_address(address, priority)
if ssl:
self.establish_ssl(server=True)
def restore_server(self):
if not hasattr(self, "_backup_server_conn"):
return
self.log("Restore original server connection: %s:%s -> %s:%s" % (
self.c.server_conn.address.host,
self.c.server_conn.address.port,
self._backup_server_conn.host,
self._backup_server_conn.port
))
self.c.del_server_connection()
self.c.server_conn = self._backup_server_conn
del self._backup_server_conn
from . import http, tcp
protocols = {
'http': dict(handler=http.HTTPHandler, flow=http.HTTPFlow),
'tcp': dict(handler=tcp.TCPHandler)
} # PyCharm type hinting behaves bad if this is a dict constructor...
def _handler(conntype, connection_handler):
if conntype in protocols:
return protocols[conntype]["handler"](connection_handler)
raise NotImplementedError # pragma: nocover
def handle_messages(conntype, connection_handler):
return _handler(conntype, connection_handler).handle_messages()
def handle_error(conntype, connection_handler, error):
return _handler(conntype, connection_handler).handle_error(error)

View File

@ -0,0 +1,21 @@
from . import http, tcp
protocols = {
'http': dict(handler=http.HTTPHandler, flow=http.HTTPFlow),
'tcp': dict(handler=tcp.TCPHandler)
} # PyCharm type hinting behaves bad if this is a dict constructor...
def _handler(conntype, connection_handler):
if conntype in protocols:
return protocols[conntype]["handler"](connection_handler)
raise NotImplementedError # pragma: nocover
def handle_messages(conntype, connection_handler):
return _handler(conntype, connection_handler).handle_messages()
def handle_error(conntype, connection_handler, error):
return _handler(conntype, connection_handler).handle_error(error)

View File

@ -1,15 +1,11 @@
import Cookie, urllib, urlparse, time, copy import Cookie, urllib, urlparse, time, copy
from email.utils import parsedate_tz, formatdate, mktime_tz from email.utils import parsedate_tz, formatdate, mktime_tz
from libmproxy.proxy.primitives import AddressPriority
from ..proxy.connection import ServerConnection
from ..proxy.primitives import ProxyError, ConnectionTypeChange
import netlib.utils import netlib.utils
from netlib import http, tcp, http_status from netlib import http, tcp, http_status
from netlib.odict import ODict, ODictCaseless from netlib.odict import ODict, ODictCaseless
from . import ProtocolHandler, KILL, TemporaryServerChangeMixin from .primitives import KILL, ProtocolHandler, TemporaryServerChangeMixin, Flow, Error
from .. import encoding, utils, filt, controller, stateobject from ..proxy.connection import ServerConnection
from .primitives import Flow, Error from .. import encoding, utils, filt, controller, stateobject, proxy
HDR_FORM_URLENCODED = "application/x-www-form-urlencoded" HDR_FORM_URLENCODED = "application/x-www-form-urlencoded"
CONTENT_MISSING = 0 CONTENT_MISSING = 0
@ -350,7 +346,7 @@ class HTTPRequest(HTTPMessage):
Raises an Exception if the request cannot be assembled. Raises an Exception if the request cannot be assembled.
""" """
if self.content == CONTENT_MISSING: if self.content == CONTENT_MISSING:
raise ProxyError(502, "Cannot assemble flow with CONTENT_MISSING") raise proxy.ProxyError(502, "Cannot assemble flow with CONTENT_MISSING")
head = self._assemble_head(form) head = self._assemble_head(form)
if self.content: if self.content:
return head + self.content return head + self.content
@ -513,7 +509,7 @@ class HTTPRequest(HTTPMessage):
self.flow.change_server((host, port), ssl=is_ssl) self.flow.change_server((host, port), ssl=is_ssl)
else: else:
# There's not live server connection, we're just changing the attributes here. # There's not live server connection, we're just changing the attributes here.
self.flow.server_conn = ServerConnection((host, port), AddressPriority.MANUALLY_CHANGED) self.flow.server_conn = ServerConnection((host, port), proxy.AddressPriority.MANUALLY_CHANGED)
self.flow.server_conn.ssl_established = is_ssl self.flow.server_conn.ssl_established = is_ssl
# If this is an absolute request, replace the attributes on the request object as well. # If this is an absolute request, replace the attributes on the request object as well.
@ -651,7 +647,7 @@ class HTTPResponse(HTTPMessage):
Raises an Exception if the request cannot be assembled. Raises an Exception if the request cannot be assembled.
""" """
if self.content == CONTENT_MISSING: if self.content == CONTENT_MISSING:
raise ProxyError(502, "Cannot assemble flow with CONTENT_MISSING") raise proxy.ProxyError(502, "Cannot assemble flow with CONTENT_MISSING")
head = self._assemble_head() head = self._assemble_head()
if self.content: if self.content:
return head + self.content return head + self.content
@ -927,7 +923,7 @@ class HTTPHandler(ProtocolHandler, TemporaryServerChangeMixin):
self.restore_server() # If the user has changed the target server on this connection, self.restore_server() # If the user has changed the target server on this connection,
# restore the original target server # restore the original target server
return True return True
except (HttpAuthenticationError, http.HttpError, ProxyError, tcp.NetLibError), e: except (HttpAuthenticationError, http.HttpError, proxy.ProxyError, tcp.NetLibError), e:
self.handle_error(e, flow) self.handle_error(e, flow)
return False return False
@ -937,7 +933,7 @@ class HTTPHandler(ProtocolHandler, TemporaryServerChangeMixin):
code = 407 code = 407
message = "Proxy Authentication Required" message = "Proxy Authentication Required"
headers = error.auth_headers headers = error.auth_headers
elif isinstance(error, (http.HttpError, ProxyError)): elif isinstance(error, (http.HttpError, proxy.ProxyError)):
code = error.code code = error.code
message = error.msg message = error.msg
elif isinstance(error, tcp.NetLibError): elif isinstance(error, tcp.NetLibError):
@ -1000,7 +996,7 @@ class HTTPHandler(ProtocolHandler, TemporaryServerChangeMixin):
self.c.log("Hook: Read answer to CONNECT request from proxy") self.c.log("Hook: Read answer to CONNECT request from proxy")
resp = HTTPResponse.from_stream(self.c.server_conn.rfile, upstream_request.method) resp = HTTPResponse.from_stream(self.c.server_conn.rfile, upstream_request.method)
if resp.code != 200: if resp.code != 200:
raise ProxyError(resp.code, raise proxy.ProxyError(resp.code,
"Cannot reestablish SSL connection with upstream proxy: \r\n" + str(resp.headers)) "Cannot reestablish SSL connection with upstream proxy: \r\n" + str(resp.headers))
self.c.log("Hook: Establish SSL with upstream proxy") self.c.log("Hook: Establish SSL with upstream proxy")
self.c.establish_ssl(server=True) self.c.establish_ssl(server=True)
@ -1028,7 +1024,7 @@ class HTTPHandler(ProtocolHandler, TemporaryServerChangeMixin):
if self.expected_form_in == "absolute": if self.expected_form_in == "absolute":
if not self.c.config.get_upstream_server: if not self.c.config.get_upstream_server:
self.c.set_server_address((request.host, request.port), AddressPriority.FROM_PROTOCOL) self.c.set_server_address((request.host, request.port), proxy.AddressPriority.FROM_PROTOCOL)
flow.server_conn = self.c.server_conn # Update server_conn attribute on the flow flow.server_conn = self.c.server_conn # Update server_conn attribute on the flow
self.c.client_conn.send( self.c.client_conn.send(
'HTTP/1.1 200 Connection established\r\n' + 'HTTP/1.1 200 Connection established\r\n' +
@ -1046,7 +1042,7 @@ class HTTPHandler(ProtocolHandler, TemporaryServerChangeMixin):
if request.scheme != "http": if request.scheme != "http":
raise http.HttpError(400, "Invalid request scheme: %s" % request.scheme) raise http.HttpError(400, "Invalid request scheme: %s" % request.scheme)
self.c.set_server_address((request.host, request.port), AddressPriority.FROM_PROTOCOL) self.c.set_server_address((request.host, request.port), proxy.AddressPriority.FROM_PROTOCOL)
flow.server_conn = self.c.server_conn # Update server_conn attribute on the flow flow.server_conn = self.c.server_conn # Update server_conn attribute on the flow
request.form_out = self.expected_form_out request.form_out = self.expected_form_out

View File

@ -1,8 +1,12 @@
from .. import stateobject, utils, version from .. import stateobject, utils, version
from ..proxy.primitives import AddressPriority
from ..proxy.connection import ClientConnection, ServerConnection from ..proxy.connection import ClientConnection, ServerConnection
import copy import copy
KILL = 0 # const for killed requests
class BackreferenceMixin(object): class BackreferenceMixin(object):
""" """
If an attribute from the _backrefattr tuple is set, If an attribute from the _backrefattr tuple is set,
@ -127,4 +131,72 @@ class Flow(stateobject.SimpleStateObject, BackreferenceMixin):
""" """
if self._backup: if self._backup:
self._load_state(self._backup) self._load_state(self._backup)
self._backup = None self._backup = None
class ProtocolHandler(object):
def __init__(self, c):
self.c = c
"""@type: libmproxy.proxy.ConnectionHandler"""
def handle_messages(self):
"""
This method gets called if a client connection has been made. Depending on the proxy settings,
a server connection might already exist as well.
"""
raise NotImplementedError # pragma: nocover
def handle_error(self, error):
"""
This method gets called should there be an uncaught exception during the connection.
This might happen outside of handle_messages, e.g. if the initial SSL handshake fails in transparent mode.
"""
raise error # pragma: nocover
class TemporaryServerChangeMixin(object):
"""
This mixin allows safe modification of the target server,
without any need to expose the ConnectionHandler to the Flow.
"""
def change_server(self, address, ssl):
if address == self.c.server_conn.address():
return
priority = AddressPriority.MANUALLY_CHANGED
if self.c.server_conn.priority > priority:
self.log("Attempt to change server address, "
"but priority is too low (is: %s, got: %s)" % (self.server_conn.priority, priority))
return
self.log("Temporarily change server connection: %s:%s -> %s:%s" % (
self.c.server_conn.address.host,
self.c.server_conn.address.port,
address.host,
address.port
))
if not hasattr(self, "_backup_server_conn"):
self._backup_server_conn = self.c.server_conn
self.c.server_conn = None
else: # This is at least the second temporary change. We can kill the current connection.
self.c.del_server_connection()
self.c.set_server_address(address, priority)
if ssl:
self.c.establish_ssl(server=True)
def restore_server(self):
if not hasattr(self, "_backup_server_conn"):
return
self.log("Restore original server connection: %s:%s -> %s:%s" % (
self.c.server_conn.address.host,
self.c.server_conn.address.port,
self._backup_server_conn.host,
self._backup_server_conn.port
))
self.c.del_server_connection()
self.c.server_conn = self._backup_server_conn
del self._backup_server_conn

View File

@ -1,6 +1,6 @@
from . import ProtocolHandler
import select, socket import select, socket
from cStringIO import StringIO from cStringIO import StringIO
from .primitives import ProtocolHandler
class TCPHandler(ProtocolHandler): class TCPHandler(ProtocolHandler):

View File

@ -0,0 +1 @@
from .primitives import *

View File

@ -3,7 +3,8 @@ from OpenSSL import SSL
from netlib import tcp from netlib import tcp
from .primitives import ProxyServerError, Log, ProxyError, ConnectionTypeChange, AddressPriority from .primitives import ProxyServerError, Log, ProxyError, ConnectionTypeChange, AddressPriority
from .connection import ClientConnection, ServerConnection from .connection import ClientConnection, ServerConnection
from .. import version, protocol from ..protocol.handle import handle_messages, handle_error
from .. import version
class DummyServer: class DummyServer:
@ -81,14 +82,14 @@ class ConnectionHandler:
while not self.close: while not self.close:
try: try:
protocol.handle_messages(self.conntype, self) handle_messages(self.conntype, self)
except ConnectionTypeChange: except ConnectionTypeChange:
self.log("Connection Type Changed: %s" % self.conntype) self.log("Connection Type Changed: %s" % self.conntype)
continue continue
# FIXME: Do we want to persist errors? # FIXME: Do we want to persist errors?
except (ProxyError, tcp.NetLibError), e: except (ProxyError, tcp.NetLibError), e:
protocol.handle_error(self.conntype, self, e) handle_error(self.conntype, self, e)
except Exception, e: except Exception, e:
self.log(e.__class__) self.log(e.__class__)
import traceback import traceback