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"])
raise FlowReadError("Incompatible serialized data version: %s"%v)
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:
# Error is due to EOF
if self.fo.tell() == off and self.fo.read() == '':

View File

@ -1,95 +1 @@
from ..proxy.primitives import AddressPriority
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)
from .primitives import *

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
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
from netlib import http, tcp, http_status
from netlib.odict import ODict, ODictCaseless
from . import ProtocolHandler, KILL, TemporaryServerChangeMixin
from .. import encoding, utils, filt, controller, stateobject
from .primitives import Flow, Error
from .primitives import KILL, ProtocolHandler, TemporaryServerChangeMixin, Flow, Error
from ..proxy.connection import ServerConnection
from .. import encoding, utils, filt, controller, stateobject, proxy
HDR_FORM_URLENCODED = "application/x-www-form-urlencoded"
CONTENT_MISSING = 0
@ -350,7 +346,7 @@ class HTTPRequest(HTTPMessage):
Raises an Exception if the request cannot be assembled.
"""
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)
if self.content:
return head + self.content
@ -513,7 +509,7 @@ class HTTPRequest(HTTPMessage):
self.flow.change_server((host, port), ssl=is_ssl)
else:
# 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
# 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.
"""
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()
if 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,
# restore the original target server
return True
except (HttpAuthenticationError, http.HttpError, ProxyError, tcp.NetLibError), e:
except (HttpAuthenticationError, http.HttpError, proxy.ProxyError, tcp.NetLibError), e:
self.handle_error(e, flow)
return False
@ -937,7 +933,7 @@ class HTTPHandler(ProtocolHandler, TemporaryServerChangeMixin):
code = 407
message = "Proxy Authentication Required"
headers = error.auth_headers
elif isinstance(error, (http.HttpError, ProxyError)):
elif isinstance(error, (http.HttpError, proxy.ProxyError)):
code = error.code
message = error.msg
elif isinstance(error, tcp.NetLibError):
@ -1000,7 +996,7 @@ class HTTPHandler(ProtocolHandler, TemporaryServerChangeMixin):
self.c.log("Hook: Read answer to CONNECT request from proxy")
resp = HTTPResponse.from_stream(self.c.server_conn.rfile, upstream_request.method)
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))
self.c.log("Hook: Establish SSL with upstream proxy")
self.c.establish_ssl(server=True)
@ -1028,7 +1024,7 @@ class HTTPHandler(ProtocolHandler, TemporaryServerChangeMixin):
if self.expected_form_in == "absolute":
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
self.c.client_conn.send(
'HTTP/1.1 200 Connection established\r\n' +
@ -1046,7 +1042,7 @@ class HTTPHandler(ProtocolHandler, TemporaryServerChangeMixin):
if request.scheme != "http":
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
request.form_out = self.expected_form_out

View File

@ -1,8 +1,12 @@
from .. import stateobject, utils, version
from ..proxy.primitives import AddressPriority
from ..proxy.connection import ClientConnection, ServerConnection
import copy
KILL = 0 # const for killed requests
class BackreferenceMixin(object):
"""
If an attribute from the _backrefattr tuple is set,
@ -128,3 +132,71 @@ class Flow(stateobject.SimpleStateObject, BackreferenceMixin):
if self._backup:
self._load_state(self._backup)
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
from cStringIO import StringIO
from .primitives import 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 .primitives import ProxyServerError, Log, ProxyError, ConnectionTypeChange, AddressPriority
from .connection import ClientConnection, ServerConnection
from .. import version, protocol
from ..protocol.handle import handle_messages, handle_error
from .. import version
class DummyServer:
@ -81,14 +82,14 @@ class ConnectionHandler:
while not self.close:
try:
protocol.handle_messages(self.conntype, self)
handle_messages(self.conntype, self)
except ConnectionTypeChange:
self.log("Connection Type Changed: %s" % self.conntype)
continue
# FIXME: Do we want to persist errors?
except (ProxyError, tcp.NetLibError), e:
protocol.handle_error(self.conntype, self, e)
handle_error(self.conntype, self, e)
except Exception, e:
self.log(e.__class__)
import traceback