def connect_application_session(server_factory,
MyApplicationSession,
component_config=None):
"""
Connect an ApplicationSession class to the given server factory.
"""
application_session_factory = ApplicationSessionFactory(component_config)
application_session_factory.session = MyApplicationSession
client_factory = WampRawSocketClientFactory(application_session_factory)
# client_factory = WampWebSocketClientFactory(application_session_factory)
server_protocol = server_factory.buildProtocol(None)
client_protocol = client_factory.buildProtocol(None)
server_transport = FakeTransport(server_protocol, True)
client_transport = FakeTransport(client_protocol, False)
pump = connect(server_protocol,
server_transport,
client_protocol,
client_transport,
debug=False)
return client_protocol._session, pump
def connectedServerAndClient(self, serverFactory, clientFactory):
"""
Set up an in-memory connection between protocols created by
C{serverFactory} and C{clientFactory}.
@return: A three-tuple. The first element is the protocol created by
C{serverFactory}. The second element is the protocol created by
C{clientFactory}. The third element is the L{IOPump} connecting
them.
"""
clientProtocol = clientFactory.buildProtocol(None)
serverProtocol = serverFactory.buildProtocol(None)
clientTransport = FakeTransport(clientProtocol,
isServer=False,
hostAddress=self.clientAddress,
peerAddress=self.serverAddress)
serverTransport = FakeTransport(serverProtocol,
isServer=True,
hostAddress=self.serverAddress,
peerAddress=self.clientAddress)
pump = connect(serverProtocol, serverTransport, clientProtocol,
clientTransport)
return serverProtocol, clientProtocol, pump
def setUp(self):
self.peer = IPv4Address('TCP', 'kafka', 9072)
self.protocol = KafkaBootstrapProtocol()
self.transport = FakeTransport(self.protocol,
isServer=False,
peerAddress=self.peer)
self.protocol.makeConnection(self.transport)
def test_connect_socks_unknown_version(self):
class BadSocksServer(Protocol):
def __init__(self):
self._buffer = b''
self._recv_stack = [
(b'\x05\x01\x00', b'\x05\xff'),
]
def dataReceived(self, data):
self._buffer += data
if len(self._recv_stack) == 0:
assert "not expecting any more data, got {}".format(repr(self._buffer))
return
expecting, to_send = self._recv_stack.pop(0)
got = self._buffer[:len(expecting)]
self._buffer = self._buffer[len(expecting):]
assert got == expecting, "wanted {} but got {}".format(repr(expecting), repr(got))
self.transport.write(to_send)
factory = socks._TorSocksFactory(u'1.2.3.4', 1234, 'CONNECT', Mock())
server_proto = BadSocksServer()
server_transport = FakeTransport(server_proto, isServer=True)
client_proto = factory.buildProtocol('ignored')
client_transport = FakeTransport(client_proto, isServer=False)
# returns IOPump
yield connect(
server_proto, server_transport,
client_proto, client_transport,
)
self.assertTrue(server_proto.transport.disconnected)
self.assertTrue(client_proto.transport.disconnected)
def test_connect_socks_illegal_packet(self):
class BadSocksServer(Protocol):
def __init__(self):
self._buffer = b''
def dataReceived(self, data):
self._buffer += data
if len(self._buffer) == 3:
assert self._buffer == b'\x05\x01\x00'
self._buffer = b''
self.transport.write(b'\x05\x01\x01')
factory = socks._TorSocksFactory(u'meejah.ca', 1234, 'CONNECT', Mock())
server_proto = BadSocksServer()
server_transport = FakeTransport(server_proto, isServer=True)
client_proto = factory.buildProtocol('ignored')
client_transport = FakeTransport(client_proto, isServer=False)
pump = yield connect(
server_proto,
server_transport,
client_proto,
client_transport,
)
self.assertTrue(server_proto.transport.disconnected)
self.assertTrue(client_proto.transport.disconnected)
pump.flush()
def test_writeSequence(self):
"""
L{FakeTransport.writeSequence} will write a sequence of L{bytes} to the
transport.
"""
a = FakeTransport(object(), False)
a.write(b"a")
a.writeSequence([b"b", b"c", b"d"])
self.assertEqual(b"".join(a.stream), b"abcd")
def getTransportAndProtocol(self, function):
"""
Construct a L{GreenletProtocol} with the given C{function} and create a
L{FakeTransport} to go along with it.
@param function: Function to pass to L{GreenletProtocol}
@rtype: Two-tuple of L{FakeTransport}, L{GreenletProtocol}.
"""
twistedTransport = FakeTransport()
factory = _GreenletFactory(function)
protocol = factory.buildProtocol(None)
twistedTransport.protocol = protocol
return twistedTransport, protocol
def test_socks_ipv6(self):
class BadSocksServer(Protocol):
def __init__(self):
self._buffer = b''
self._recv_stack = [
(b'\x05\x01\x00', b'\x05\x00'),
(b'\x05\x01\x00\x04\x20\x02\x44\x93\x04\xd2',
b'\x05\x00\x00\x04%s\xbe\xef' % (b'\x00' * 16)),
]
def dataReceived(self, data):
self._buffer += data
if len(self._recv_stack) == 0:
assert "not expecting any more data, got {}".format(
repr(self._buffer))
return
expecting, to_send = self._recv_stack.pop(0)
got = self._buffer[:len(expecting)]
self._buffer = self._buffer[len(expecting):]
assert got == expecting, "wanted {} but got {}".format(
repr(expecting), repr(got))
self.transport.write(to_send)
factory = socks._TorSocksFactory(
u'2002:4493:5105::a299:9bff:fe0e:4471', 1234, 'CONNECT', Mock())
server_proto = BadSocksServer()
expected_address = object()
server_transport = FakeTransport(server_proto, isServer=True)
client_proto = factory.buildProtocol(u'ignored')
client_transport = FakeTransport(client_proto,
isServer=False,
hostAddress=expected_address)
pump = yield connect(
server_proto,
server_transport,
client_proto,
client_transport,
)
# should be relaying now, try sending some datas
client_proto.transport.write(b'abcdef')
addr = yield factory._get_address()
# FIXME how shall we test for IPv6-ness?
assert addr is expected_address
pump.flush()
self.assertEqual(b'abcdef', server_proto._buffer)
def connect(self, factory):
if self.pump is not None:
raise Exception("SingleUseMemoryEndpoint was already used")
try:
protocol = factory.buildProtocol(MemoryAddress())
except:
return fail()
else:
self.pump = connect(self._server,
FakeTransport(self._server, isServer=True),
protocol,
FakeTransport(protocol, isServer=False))
return succeed(protocol)
def unconnected_proxyserver(mocker):
mocker.patch("twisted.test.iosim.FakeTransport.startTLS")
mocker.patch("pappyproxy.proxy.load_certs_from_dir",
new=mock_generate_cert)
factory = ProxyServerFactory()
protocol = factory.buildProtocol(('127.0.0.1', 0))
protocol.makeConnection(FakeTransport(protocol, True))
return protocol
def test_connect_socks_unknown_reply_code(self):
class BadSocksServer(Protocol):
def __init__(self):
self._buffer = b''
self._recv_stack = [
(b'\x05\x01\x00', b'\x05\x00'),
# the \xff is an invalid reply-code
(b'\x05\x01\x00\x01\x01\x02\x03\x04\x04\xd2',
b'\x05\xff\x00\x04\x01\x01\x01\x01'),
]
def dataReceived(self, data):
self._buffer += data
if len(self._recv_stack) == 0:
assert "not expecting any more data, got {}".format(
repr(self._buffer))
return
expecting, to_send = self._recv_stack.pop(0)
got = self._buffer[:len(expecting)]
self._buffer = self._buffer[len(expecting):]
assert got == expecting, "wanted {} but got {}".format(
repr(expecting), repr(got))
self.transport.write(to_send)
factory = socks._TorSocksFactory(u'1.2.3.4', 1234, 'CONNECT', Mock())
server_proto = BadSocksServer()
server_transport = FakeTransport(server_proto, isServer=True)
client_proto = factory.buildProtocol('ignored')
client_transport = FakeTransport(client_proto, isServer=False)
d = client_proto._machine.when_done()
# returns IOPump
yield connect(
server_proto,
server_transport,
client_proto,
client_transport,
)
with self.assertRaises(Exception) as ctx:
yield d
self.assertIn('Unknown SOCKS error-code', str(ctx.exception))
def test_socks_relay_data(self):
class BadSocksServer(Protocol):
def __init__(self):
self._buffer = b''
self._recv_stack = [
(b'\x05\x01\x00', b'\x05\x00'),
(b'\x05\x01\x00\x01\x01\x02\x03\x04\x04\xd2',
b'\x05\x00\x00\x01\x01\x02\x03\x04\x12\x34'),
]
def dataReceived(self, data):
self._buffer += data
if len(self._recv_stack) == 0:
assert "not expecting any more data, got {}".format(
repr(self._buffer))
return
expecting, to_send = self._recv_stack.pop(0)
got = self._buffer[:len(expecting)]
self._buffer = self._buffer[len(expecting):]
assert got == expecting, "wanted {} but got {}".format(
repr(expecting), repr(got))
self.transport.write(to_send)
factory = socks._TorSocksFactory(u'1.2.3.4', 1234, 'CONNECT', Mock())
server_proto = BadSocksServer()
server_transport = FakeTransport(server_proto, isServer=True)
client_proto = factory.buildProtocol('ignored')
client_transport = FakeTransport(client_proto, isServer=False)
pump = yield connect(
server_proto,
server_transport,
client_proto,
client_transport,
)
# should be relaying now, try sending some datas
client_proto.transport.write(b'abcdef')
pump.flush()
self.assertEqual(b'abcdef', server_proto._buffer)
def proxyserver(mocker):
mocker.patch("twisted.test.iosim.FakeTransport.startTLS")
mocker.patch("pappyproxy.proxy.load_certs_from_dir",
new=mock_generate_cert)
factory = ProxyServerFactory()
protocol = factory.buildProtocol(('127.0.0.1', 0))
protocol.makeConnection(FakeTransport(protocol, True))
protocol.lineReceived('CONNECT https://www.AAAA.BBBB:443 HTTP/1.1')
protocol.lineReceived('')
protocol.transport.getOutBuffer()
return protocol
def gen_connection(send_data,
new_req=False,
new_rsp=False,
drop_req=False,
drop_rsp=False):
factory = ProxyClientFactory(http.Request(send_data))
macro = gen_mangle_macro(new_req, new_rsp, drop_req, drop_rsp)
factory.intercepting_macros['pappy_mangle'] = macro
protocol = factory.buildProtocol(None)
tr = FakeTransport(protocol, True)
protocol.makeConnection(tr)
sent = yield protocol.data_defer
print sent
defer.returnValue((protocol, sent, factory.data_defer))
class KafkaBootstrapProtocolTests(SynchronousTestCase):
"""Test `afkak._protocol.KafkaBootstrapProtocol`
:ivar peer: Peer IAddress, an IPv4 address
:ivar protocol: `KafkaBootstrapProtocol` object, connected in setUp
:ivar transport: `FakeTransport` object associated with the protocol,
connected in setUp
"""
def setUp(self):
self.peer = IPv4Address('TCP', 'kafka', 9072)
self.protocol = KafkaBootstrapProtocol()
self.transport = FakeTransport(self.protocol, isServer=False, peerAddress=self.peer)
self.protocol.makeConnection(self.transport)
def test_one_request(self):
"""
Happy path: a request is made and a response received.
"""
correlation_id = b'corr'
client_request = b'api ' + correlation_id + b'corr more stuff'
d = self.protocol.request(client_request)
self.assertNoResult(d)
# The request was written to the server.
server_request = self.transport.getOutBuffer()
self.assertEqual(b'\0\0\0\x17' + client_request, server_request)
self.transport.bufferReceived(b'\0\0\0\x05' + correlation_id + b'y')
self.assertEqual(correlation_id + b'y', self.successResultOf(d))
def test_disconnected(self):
"""
Pending and future requests fail when the connection goes away.
"""
d = self.protocol.request(b'api corr stuff')
self.assertNoResult(d)
self.transport.disconnectReason = ConnectionLost('Bye.')
self.transport.reportDisconnect()
self.failureResultOf(d, ConnectionLost)
self.failureResultOf(self.protocol.request(b'api corr more'), ConnectionLost)
def test_unknown_correlation_id(self):
"""
A warning is logged and the connection dropped when a response with an
unknown correlation ID is received.
"""
events = EventLoggingObserver.createWithCleanup(self, globalLogPublisher)
self.transport.bufferReceived(b'\0\0\0\x101234 more stuff..')
self.assertTrue(self.transport.disconnecting)
[event] = events
self.assertEqual(LogLevel.warn, event['log_level'])
self.assertEqual(self.peer, event['peer'])
self.assertEqual(b'1234', event['correlation_id'])
def test_oversized_response(self):
"""
An oversized response from the server prompts disconnection.
"""
d = self.protocol.request(b'api corr blah blah')
self.transport.bufferReceived(b'\xff\xff\xff\xff') # 2**32 - 1, way too large.
self.assertTrue(self.transport.disconnecting)
self.assertNoResult(d) # Will fail when the disconnect completes.
def __init__(self, protocol, isServer, q2qhost, q2qpeer):
FakeTransport.__init__(self, protocol, isServer)
self.q2qhost = q2qhost
self.q2qpeer = q2qpeer
class KafkaBootstrapProtocolTests(SynchronousTestCase):
"""Test `afkak._protocol.KafkaBootstrapProtocol`
:ivar peer: Peer IAddress, an IPv4 address
:ivar protocol: `KafkaBootstrapProtocol` object, connected in setUp
:ivar transport: `FakeTransport` object associated with the protocol,
connected in setUp
"""
def setUp(self):
self.peer = IPv4Address('TCP', 'kafka', 9072)
self.protocol = KafkaBootstrapProtocol()
self.transport = FakeTransport(self.protocol,
isServer=False,
peerAddress=self.peer)
self.protocol.makeConnection(self.transport)
def test_one_request(self):
"""
Happy path: a request is made and a response received.
"""
correlation_id = b'corr'
client_request = b'api ' + correlation_id + b'corr more stuff'
d = self.protocol.request(client_request)
self.assertNoResult(d)
# The request was written to the server.
server_request = self.transport.getOutBuffer()
self.assertEqual(b'\0\0\0\x17' + client_request, server_request)
self.transport.bufferReceived(b'\0\0\0\x05' + correlation_id + b'y')
self.assertEqual(correlation_id + b'y', self.successResultOf(d))
def test_disconnected(self):
"""
Pending and future requests fail when the connection goes away.
"""
d = self.protocol.request(b'api corr stuff')
self.assertNoResult(d)
self.transport.disconnectReason = ConnectionLost('Bye.')
self.transport.reportDisconnect()
self.failureResultOf(d, ConnectionLost)
self.failureResultOf(self.protocol.request(b'api corr more'),
ConnectionLost)
def test_unknown_correlation_id(self):
"""
A warning is logged and the connection dropped when a response with an
unknown correlation ID is received.
"""
events = EventLoggingObserver.createWithCleanup(
self, globalLogPublisher)
self.transport.bufferReceived(b'\0\0\0\x101234 more stuff..')
self.assertTrue(self.transport.disconnecting)
[event] = events
self.assertEqual(LogLevel.warn, event['log_level'])
self.assertEqual(self.peer, event['peer'])
self.assertEqual(b'1234', event['correlation_id'])
def test_oversized_response(self):
"""
An oversized response from the server prompts disconnection.
"""
d = self.protocol.request(b'api corr blah blah')
self.transport.bufferReceived(
b'\xff\xff\xff\xff') # 2**32 - 1, way too large.
self.assertTrue(self.transport.disconnecting)
self.assertNoResult(d) # Will fail when the disconnect completes.
def __init__(self, protocol, isServer, q2qhost, q2qpeer):
FakeTransport.__init__(self, protocol, isServer)
self.q2qhost = q2qhost
self.q2qpeer = q2qpeer
def setUp(self):
self.peer = IPv4Address('TCP', 'kafka', 9072)
self.protocol = KafkaBootstrapProtocol()
self.transport = FakeTransport(self.protocol, isServer=False, peerAddress=self.peer)
self.protocol.makeConnection(self.transport)
def __init__(self, q2qhost, q2qpeer):
FakeTransport.__init__(self)
self.q2qhost = q2qhost
self.q2qpeer = q2qpeer
