# Copyright (c) Twisted Matrix Laboratories.
# See LICENSE for details.

"""
Test cases for twisted.protocols package.
"""

from twisted.trial import unittest
from twisted.protocols import wire, portforward
from twisted.internet import reactor, defer, address, protocol
from twisted.test import proto_helpers


class WireTestCase(unittest.TestCase):
"""
Test wire protocols.
"""

def test_echo(self):
"""
Test wire.Echo protocol: send some data and check it send it back.
"""
t = proto_helpers.StringTransport()
a = wire.Echo()
a.makeConnection(t)
a.dataReceived("hello")
a.dataReceived("world")
a.dataReceived("how")
a.dataReceived("are")
a.dataReceived("you")
self.assertEqual(t.value(), "helloworldhowareyou")


def test_who(self):
"""
Test wire.Who protocol.
"""
t = proto_helpers.StringTransport()
a = wire.Who()
a.makeConnection(t)
self.assertEqual(t.value(), "root\r\n")


def test_QOTD(self):
"""
Test wire.QOTD protocol.
"""
t = proto_helpers.StringTransport()
a = wire.QOTD()
a.makeConnection(t)
self.assertEqual(t.value(),
"An apple a day keeps the doctor away.\r\n")


def test_discard(self):
"""
Test wire.Discard protocol.
"""
t = proto_helpers.StringTransport()
a = wire.Discard()
a.makeConnection(t)
a.dataReceived("hello")
a.dataReceived("world")
a.dataReceived("how")
a.dataReceived("are")
a.dataReceived("you")
self.assertEqual(t.value(), "")



class TestableProxyClientFactory(portforward.ProxyClientFactory):
"""
Test proxy client factory that keeps the last created protocol instance.

@ivar protoInstance: the last instance of the protocol.
@type protoInstance: L{portforward.ProxyClient}
"""

def buildProtocol(self, addr):
"""
Create the protocol instance and keeps track of it.
"""
proto = portforward.ProxyClientFactory.buildProtocol(self, addr)
self.protoInstance = proto
return proto



class TestableProxyFactory(portforward.ProxyFactory):
"""
Test proxy factory that keeps the last created protocol instance.

@ivar protoInstance: the last instance of the protocol.
@type protoInstance: L{portforward.ProxyServer}

@ivar clientFactoryInstance: client factory used by C{protoInstance} to
create forward connections.
@type clientFactoryInstance: L{TestableProxyClientFactory}
"""

def buildProtocol(self, addr):
"""
Create the protocol instance, keeps track of it, and makes it use
C{clientFactoryInstance} as client factory.
"""
proto = portforward.ProxyFactory.buildProtocol(self, addr)
self.clientFactoryInstance = TestableProxyClientFactory()
# Force the use of this specific instance
proto.clientProtocolFactory = lambda: self.clientFactoryInstance
self.protoInstance = proto
return proto



class Portforwarding(unittest.TestCase):
"""
Test port forwarding.
"""

def setUp(self):
self.serverProtocol = wire.Echo()
self.clientProtocol = protocol.Protocol()
self.openPorts = []


def tearDown(self):
try:
self.proxyServerFactory.protoInstance.transport.loseConnection()
except AttributeError:
pass
try:
pi = self.proxyServerFactory.clientFactoryInstance.protoInstance
pi.transport.loseConnection()
except AttributeError:
pass
try:
self.clientProtocol.transport.loseConnection()
except AttributeError:
pass
try:
self.serverProtocol.transport.loseConnection()
except AttributeError:
pass
return defer.gatherResults(
[defer.maybeDeferred(p.stopListening) for p in self.openPorts])


def test_portforward(self):
"""
Test port forwarding through Echo protocol.
"""
realServerFactory = protocol.ServerFactory()
realServerFactory.protocol = lambda: self.serverProtocol
realServerPort = reactor.listenTCP(0, realServerFactory,
interface='127.0.0.1')
self.openPorts.append(realServerPort)
self.proxyServerFactory = TestableProxyFactory('127.0.0.1',
realServerPort.getHost().port)
proxyServerPort = reactor.listenTCP(0, self.proxyServerFactory,
interface='127.0.0.1')
self.openPorts.append(proxyServerPort)

nBytes = 1000
received = []
d = defer.Deferred()

def testDataReceived(data):
received.extend(data)
if len(received) >= nBytes:
self.assertEqual(''.join(received), 'x' * nBytes)
d.callback(None)

self.clientProtocol.dataReceived = testDataReceived

def testConnectionMade():
self.clientProtocol.transport.write('x' * nBytes)

self.clientProtocol.connectionMade = testConnectionMade

clientFactory = protocol.ClientFactory()
clientFactory.protocol = lambda: self.clientProtocol

reactor.connectTCP(
'127.0.0.1', proxyServerPort.getHost().port, clientFactory)

return d


def test_registerProducers(self):
"""
The proxy client registers itself as a producer of the proxy server and
vice versa.
"""
# create a ProxyServer instance
addr = address.IPv4Address('TCP', '127.0.0.1', 0)
server = portforward.ProxyFactory('127.0.0.1', 0).buildProtocol(addr)

# set the reactor for this test
reactor = proto_helpers.MemoryReactor()
server.reactor = reactor

# make the connection
serverTransport = proto_helpers.StringTransport()
server.makeConnection(serverTransport)

# check that the ProxyClientFactory is connecting to the backend
self.assertEqual(len(reactor.tcpClients), 1)
# get the factory instance and check it's the one we expect
host, port, clientFactory, timeout, _ = reactor.tcpClients[0]
self.assertIsInstance(clientFactory, portforward.ProxyClientFactory)

# Connect it
client = clientFactory.buildProtocol(addr)
clientTransport = proto_helpers.StringTransport()
client.makeConnection(clientTransport)

# check that the producers are registered
self.assertIdentical(clientTransport.producer, serverTransport)
self.assertIdentical(serverTransport.producer, clientTransport)
# check the streaming attribute in both transports
self.assertTrue(clientTransport.streaming)
self.assertTrue(serverTransport.streaming)



class StringTransportTestCase(unittest.TestCase):
"""
Test L{proto_helpers.StringTransport} helper behaviour.
"""

def test_noUnicode(self):
"""
Test that L{proto_helpers.StringTransport} doesn't accept unicode data.
"""
s = proto_helpers.StringTransport()
self.assertRaises(TypeError, s.write, u'foo')