import os

from tempfile import NamedTemporaryFile

from twisted.trial.unittest import TestCase
from twisted.conch.manhole_ssh import ConchFactory
from twisted.conch.ssh import (
    transport, userauth, connection, channel, session)
from twisted.internet import defer, protocol, reactor
from twisted.internet.defer import inlineCallbacks

from vumi.middleware.manhole import ManholeMiddleware


# these are shipped along with Twisted
private_key = ConchFactory.privateKeys['ssh-rsa']
public_key = ConchFactory.publicKeys['ssh-rsa']


class DummyWorker(object):
    pass


class ClientTransport(transport.SSHClientTransport):

    def verifyHostKey(self, pub_key, fingerprint):
        return defer.succeed(1)

    def connectionSecure(self):
        return self.requestService(ClientUserAuth(
            os.getlogin(), ClientConnection(self.factory.channelConnected)))


class ClientUserAuth(userauth.SSHUserAuthClient):

    def getPassword(self, prompt=None):
        # Not doing password based auth
        return

    def getPublicKey(self):
        return public_key.blob()

    def getPrivateKey(self):
        return defer.succeed(private_key.keyObject)


class ClientConnection(connection.SSHConnection):

    def __init__(self, channel_connected):
        connection.SSHConnection.__init__(self)
        self._channel_connected = channel_connected

    def serviceStarted(self):
        channel = ClientChannel(self._channel_connected,
                                conn=self)
        self.openChannel(channel)


class ClientChannel(channel.SSHChannel):

    name = 'session'

    def __init__(self, channel_connected, *args, **kwargs):
        channel.SSHChannel.__init__(self, *args, **kwargs)
        self._channel_connected = channel_connected
        self.buffer = u''
        self.queue = defer.DeferredQueue()

    def channelOpen(self, data):
        self._channel_connected.callback(self)

    def dataReceived(self, data):
        self.buffer += data
        lines = self.buffer.split('\r\n')
        for line in lines[:-1]:
            self.queue.put(line)
        self.buffer = lines[-1]


class ManholeMiddlewareTestCase(TestCase):

    @inlineCallbacks
    def setUp(self):

        self.private_key_file = NamedTemporaryFile(mode='w')
        self.private_key_file.write(private_key.toString('OPENSSH'))
        self.private_key_file.flush()

        self.pub_key_file = NamedTemporaryFile(mode='w')
        self.pub_key_file.write(public_key.toString('OPENSSH'))
        self.pub_key_file.flush()

        self._middlewares = []

        mw = self.get_middleware({
            'authorized_keys': [self.pub_key_file.name]
        })
        host = mw.socket.getHost()

        factory = protocol.ClientFactory()
        factory.protocol = ClientTransport

        wait_for_channel = defer.Deferred()

        factory.channelConnected = wait_for_channel
        self.socket = reactor.connectTCP(host.host, host.port, factory)

        self.channel = yield self.open_shell(wait_for_channel)

    @inlineCallbacks
    def open_shell(self, wait_for_channel):
        channel = yield wait_for_channel
        conn = channel.conn
        term = session.packRequest_pty_req("xterm-mono", (0, 0, 0, 0), '')
        yield conn.sendRequest(channel, 'pty-req', term, wantReply=1)
        yield conn.sendRequest(channel, 'shell', '', wantReply=1)
        defer.returnValue(channel)

    def tearDown(self):
        for mw in self._middlewares:
            mw.teardown_middleware()
        self.socket.disconnect()

    def get_middleware(self, config={}):
        config = dict({
            'port': '0',
        }, **config)

        worker = DummyWorker()
        worker.transport_name = 'foo'

        mw = ManholeMiddleware("test_manhole_mw", config, worker)
        mw.setup_middleware()
        self._middlewares.append(mw)
        return mw

    @inlineCallbacks
    def test_mw(self):
        self.channel.write('print worker.transport_name\n')
        # read the echoed line we sent first
        yield self.channel.queue.get()
        # next is the server response
        received_line = yield self.channel.queue.get()
        self.assertEqual(received_line, 'foo')

@@ -2152,37 +2152,16 @@
 Case):%0A%0A
-    @inlineCallbacks%0A
     def 
@@ -2529,16 +2529,50 @@
 res = %5B%5D
+%0A        self._client_sockets = %5B%5D
 %0A%0A      
@@ -2569,24 +2569,29 @@
 %5B%5D%0A%0A        
+self.
 mw = self.ge
@@ -2671,16 +2671,76 @@
       %7D)
+%0A%0A    @inlineCallbacks%0A    def open_shell(self, middleware):
 %0A       
@@ -2748,17 +2748,25 @@
 host = m
-w
+iddleware
 .socket.
@@ -2771,25 +2771,24 @@
 t.getHost()%0A
-%0A
         fact
@@ -2861,25 +2861,24 @@
 ansport%0A
-%0A
         
 wait_for
@@ -2869,32 +2869,40 @@
         
-wait_for_channel
+factory.channelConnected
  = defer
@@ -2926,65 +2926,8 @@
     
-factory.channelConnected = wait_for_channel%0A        self.
 sock
@@ -2986,29 +2986,24 @@
 y)%0A%0A        
-self.
 channel = yi
@@ -3010,148 +3010,32 @@
 eld 
-self.open_shell(wait_for_channel)%0A%0A    @inlineCallbacks%0A    def open_shell(self, wait_for_channel):%0A        channel = yield wait_for_channel
+factory.channelConnected
 %0A   
@@ -3262,32 +3262,76 @@
 ', wantReply=1)%0A
+        self._client_sockets.append(socket)%0A
         defer.re
@@ -3366,32 +3366,109 @@
 tearDown(self):%0A
+        for socket in self._client_sockets:%0A            socket.disconnect()%0A%0A
         for mw i
@@ -3528,41 +3528,8 @@
 re()
-%0A        self.socket.disconnect()
 %0A%0A  
@@ -3913,26 +3913,66 @@
        s
-elf.channe
+hell = yield self.open_shell(self.mw)%0A        shel
 l.write(
@@ -4060,34 +4060,27 @@
      yield s
+h
 el
-f.channe
 l.queue.get(
@@ -4154,18 +4154,11 @@
 ld s
+h
 el
-f.channe
 l.qu

test cleanup