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.exceptions import HttpException
from netlib.http import Headers, http1
from netlib.utils import clean_bin
from . import controller, tnetstring, filt, script, version, flow_format_compat
from .onboarding import app
from .proxy.config import HostMatcher
from .protocol.http_replay import RequestReplayThread
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
@ -893,17 +892,6 @@ class FlowMaster(controller.ServerMaster):
self.handle_response(f)
if f.error:
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:
raise NotImplementedError()
@ -1091,39 +1079,9 @@ class FlowMaster(controller.ServerMaster):
self.add_event('"{}" reloaded.'.format(s.filename), 'info')
return ok
def handle_tcp_open(self, flow):
self.state.add_flow(flow)
self.run_script_hook("tcp_open", flow)
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 handle_tcp_message(self, m):
self.run_script_hook("tcp_message", m)
m.reply()
def shutdown(self):
super(FlowMaster, self).shutdown()

View File

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

View File

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

View File

@ -191,6 +191,12 @@ class HTTPFlow(Flow):
response=HTTPResponse
)
@classmethod
def from_state(cls, state):
f = cls(None, None)
f.set_state(state)
return f
def __repr__(self):
s = "<HTTPFlow"
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.utils import clean_bin
from ..exceptions import ProtocolException
from ..models import Error
from ..models.tcp import TCPFlow, TCPMessage
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):
chunk_size = 4096
def __init__(self, ctx, ignore=False):
self.ignore = ignore
def __init__(self, ctx, logging=True):
self.logging = logging
super(RawTCPLayer, self).__init__(ctx)
def __call__(self):
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))
client = self.client_conn.connection
@ -48,24 +51,38 @@ class RawTCPLayer(Layer):
if isinstance(conn, SSL.Connection):
# We can't half-close a connection, so we just close everything here.
# Sockets will be cleaned up on a higher level.
break
return
else:
dst.shutdown(socket.SHUT_WR)
if len(conns) == 0:
break
return
continue
tcp_message = TCPMessage(dst == server, buf[:size].tobytes())
if not self.ignore:
flow.messages.append(tcp_message)
self.channel.ask("tcp_message", flow)
dst.sendall(tcp_message.content)
tcp_message = TcpMessage(
self.client_conn, self.server_conn,
self.client_conn if dst == server else self.server_conn,
self.server_conn if dst == server else self.client_conn,
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:
if not self.ignore:
flow.error = Error("TCP connection closed unexpectedly: {}".format(repr(e)))
self.channel.tell("tcp_error", flow)
finally:
if not self.ignore:
self.channel.tell("tcp_close", flow)
six.reraise(
ProtocolException,
ProtocolException("TCP connection closed unexpectedly: {}".format(repr(e))),
sys.exc_info()[2]
)

View File

@ -65,7 +65,7 @@ class RootContext(object):
else:
ignore = self.config.check_ignore((client_hello.client_sni, 443))
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.
# An inline script may upgrade from http to https,

View File

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