Revert "fix tests"

This reverts commit f315dc1eb9.
This commit is contained in:
Maximilian Hils 2016-05-10 12:18:05 -06:00
parent f315dc1eb9
commit bb4f9611f5
7 changed files with 53 additions and 84 deletions

View File

@ -16,13 +16,12 @@ import re
from netlib import wsgi from netlib import wsgi
from netlib.exceptions import HttpException from netlib.exceptions import HttpException
from netlib.http import Headers, http1 from netlib.http import Headers, http1
from netlib.utils import clean_bin
from . import controller, tnetstring, filt, script, version, flow_format_compat from . import controller, tnetstring, filt, script, version, flow_format_compat
from .onboarding import app from .onboarding import app
from .proxy.config import HostMatcher from .proxy.config import HostMatcher
from .protocol.http_replay import RequestReplayThread from .protocol.http_replay import RequestReplayThread
from .exceptions import Kill, FlowReadException from .exceptions import Kill, FlowReadException
from .models import ClientConnection, ServerConnection, HTTPFlow, HTTPRequest, FLOW_TYPES, TCPFlow from .models import ClientConnection, ServerConnection, HTTPFlow, HTTPRequest, FLOW_TYPES
from collections import defaultdict from collections import defaultdict
@ -893,17 +892,6 @@ class FlowMaster(controller.ServerMaster):
self.handle_response(f) self.handle_response(f)
if f.error: if f.error:
self.handle_error(f) self.handle_error(f)
elif isinstance(f, TCPFlow):
messages = f.messages
f.messages = []
f.reply = controller.DummyReply()
self.handle_tcp_open(f)
while messages:
f.messages.append(messages.pop(0))
self.handle_tcp_message(f)
if f.error:
self.handle_tcp_error(f)
self.handle_tcp_close(f)
else: else:
raise NotImplementedError() raise NotImplementedError()
@ -1091,39 +1079,9 @@ class FlowMaster(controller.ServerMaster):
self.add_event('"{}" reloaded.'.format(s.filename), 'info') self.add_event('"{}" reloaded.'.format(s.filename), 'info')
return ok return ok
def handle_tcp_open(self, flow): def handle_tcp_message(self, m):
self.state.add_flow(flow) self.run_script_hook("tcp_message", m)
self.run_script_hook("tcp_open", flow) m.reply()
flow.reply()
def handle_tcp_message(self, flow):
self.run_script_hook("tcp_message", flow)
message = flow.messages[-1]
direction = "->" if message.from_client else "<-"
self.add_event("{client} {direction} tcp {direction} {server}".format(
client=repr(flow.client_conn.address),
server=repr(flow.server_conn.address),
direction=direction,
), "info")
self.add_event(clean_bin(message.content), "debug")
flow.reply()
def handle_tcp_error(self, flow):
if self.stream:
self.stream.add(flow)
self.add_event("Error in TCP connection to {}: {}".format(
repr(flow.server_conn.address),
flow.error
), "info")
self.run_script_hook("tcp_error", flow)
flow.reply()
def handle_tcp_close(self, flow):
self.state.delete_flow(flow)
if self.stream:
self.stream.add(flow)
self.run_script_hook("tcp_close", flow)
flow.reply()
def shutdown(self): def shutdown(self):
super(FlowMaster, self).shutdown() super(FlowMaster, self).shutdown()

View File

@ -7,11 +7,9 @@ from .http import (
from netlib.http import decoded from netlib.http import decoded
from .connections import ClientConnection, ServerConnection from .connections import ClientConnection, ServerConnection
from .flow import Flow, Error from .flow import Flow, Error
from .tcp import TCPFlow
FLOW_TYPES = dict( FLOW_TYPES = dict(
http=HTTPFlow, http=HTTPFlow
tcp=TCPFlow,
) )
__all__ = [ __all__ = [
@ -20,6 +18,5 @@ __all__ = [
"make_connect_response", "expect_continue_response", "make_connect_response", "expect_continue_response",
"ClientConnection", "ServerConnection", "ClientConnection", "ServerConnection",
"Flow", "Error", "Flow", "Error",
"TCPFlow"
"FLOW_TYPES" "FLOW_TYPES"
] ]

View File

@ -40,9 +40,6 @@ class Error(stateobject.StateObject):
def __str__(self): def __str__(self):
return self.msg return self.msg
def __repr__(self):
return self.msg
@classmethod @classmethod
def from_state(cls, state): def from_state(cls, state):
# the default implementation assumes an empty constructor. Override # the default implementation assumes an empty constructor. Override
@ -102,12 +99,6 @@ class Flow(stateobject.StateObject):
self._backup = state.pop("backup") self._backup = state.pop("backup")
super(Flow, self).set_state(state) super(Flow, self).set_state(state)
@classmethod
def from_state(cls, state):
f = cls(None, None)
f.set_state(state)
return f
def copy(self): def copy(self):
f = copy.copy(self) f = copy.copy(self)

View File

@ -191,6 +191,12 @@ class HTTPFlow(Flow):
response=HTTPResponse response=HTTPResponse
) )
@classmethod
def from_state(cls, state):
f = cls(None, None)
f.set_state(state)
return f
def __repr__(self): def __repr__(self):
s = "<HTTPFlow" s = "<HTTPFlow"
for a in ("request", "response", "error", "client_conn", "server_conn"): for a in ("request", "response", "error", "client_conn", "server_conn"):

View File

@ -9,26 +9,29 @@ from netlib.exceptions import TcpException
from netlib.tcp import ssl_read_select from netlib.tcp import ssl_read_select
from netlib.utils import clean_bin from netlib.utils import clean_bin
from ..exceptions import ProtocolException from ..exceptions import ProtocolException
from ..models import Error
from ..models.tcp import TCPFlow, TCPMessage
from .base import Layer 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): class RawTCPLayer(Layer):
chunk_size = 4096 chunk_size = 4096
def __init__(self, ctx, ignore=False): def __init__(self, ctx, logging=True):
self.ignore = ignore self.logging = logging
super(RawTCPLayer, self).__init__(ctx) super(RawTCPLayer, self).__init__(ctx)
def __call__(self): def __call__(self):
self.connect() self.connect()
if not self.ignore:
flow = TCPFlow(self.client_conn, self.server_conn, self)
self.channel.ask("tcp_open", flow)
buf = memoryview(bytearray(self.chunk_size)) buf = memoryview(bytearray(self.chunk_size))
client = self.client_conn.connection client = self.client_conn.connection
@ -48,24 +51,38 @@ class RawTCPLayer(Layer):
if isinstance(conn, SSL.Connection): if isinstance(conn, SSL.Connection):
# We can't half-close a connection, so we just close everything here. # We can't half-close a connection, so we just close everything here.
# Sockets will be cleaned up on a higher level. # Sockets will be cleaned up on a higher level.
break return
else: else:
dst.shutdown(socket.SHUT_WR) dst.shutdown(socket.SHUT_WR)
if len(conns) == 0: if len(conns) == 0:
break return
continue continue
tcp_message = TCPMessage(dst == server, buf[:size].tobytes()) tcp_message = TcpMessage(
if not self.ignore: self.client_conn, self.server_conn,
flow.messages.append(tcp_message) self.client_conn if dst == server else self.server_conn,
self.channel.ask("tcp_message", flow) self.server_conn if dst == server else self.client_conn,
dst.sendall(tcp_message.content) 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,
# hence the "weird" direction string.
if dst == server:
direction = "-> tcp -> {}".format(repr(self.server_conn.address))
else:
direction = "<- tcp <- {}".format(repr(self.server_conn.address))
data = clean_bin(tcp_message.message)
self.log(
"{}\r\n{}".format(direction, data),
"info"
)
except (socket.error, TcpException, SSL.Error) as e: except (socket.error, TcpException, SSL.Error) as e:
if not self.ignore: six.reraise(
flow.error = Error("TCP connection closed unexpectedly: {}".format(repr(e))) ProtocolException,
self.channel.tell("tcp_error", flow) ProtocolException("TCP connection closed unexpectedly: {}".format(repr(e))),
finally: sys.exc_info()[2]
if not self.ignore: )
self.channel.tell("tcp_close", flow)

View File

@ -65,7 +65,7 @@ class RootContext(object):
else: else:
ignore = self.config.check_ignore((client_hello.client_sni, 443)) ignore = self.config.check_ignore((client_hello.client_sni, 443))
if ignore: if ignore:
return RawTCPLayer(top_layer, ignore=True) return RawTCPLayer(top_layer, logging=False)
# 2. Always insert a TLS layer, even if there's neither client nor server tls. # 2. Always insert a TLS layer, even if there's neither client nor server tls.
# An inline script may upgrade from http to https, # An inline script may upgrade from http to https,

View File

@ -108,7 +108,7 @@ class TestRequestUtils(object):
request.url = "not-a-url" request.url = "not-a-url"
def test_url_options(self): def test_url_options(self):
request = treq(method=b"OPTIONS", path=b"*") request = treq(method="OPTIONS", path="*")
assert request.url == "http://address:22" assert request.url == "http://address:22"
def test_url_authority(self): def test_url_authority(self):
@ -149,7 +149,7 @@ class TestRequestUtils(object):
assert request.pretty_url == "http://address:22/path" assert request.pretty_url == "http://address:22/path"
def test_pretty_url_options(self): def test_pretty_url_options(self):
request = treq(method=b"OPTIONS", path=b"*") request = treq(method="OPTIONS", path="*")
assert request.pretty_url == "http://address:22" assert request.pretty_url == "http://address:22"
def test_pretty_url_authority(self): def test_pretty_url_authority(self):