← Back to branch summary

~cbehrens/nova/lp844160-build-works-with-zones

« back to all changes in this revision

Viewing changes to vendor/Twisted-10.0.0/twisted/conch/test/test_conch.py

  • Committer: Jesse Andrews
  • Date: 2010-05-28 06:05:26 UTC
  • Revision ID: git-v1:bf6e6e718cdc7488e2da87b21e258ccc065fe499
initial commit

Show diffs side-by-side

added added

removed removed

Lines of Context:
vendor/Twisted-10.0.0/twisted/conch/test/test_conch.py
 
1
# -*- test-case-name: twisted.conch.test.test_conch -*-
 
2
# Copyright (c) 2001-2008 Twisted Matrix Laboratories.
 
3
# See LICENSE for details.
 
4
 
 
5
import os, sys, socket
 
6
 
 
7
from twisted.cred import portal
 
8
from twisted.internet import reactor, defer, protocol
 
9
from twisted.internet.error import ProcessExitedAlready
 
10
from twisted.python import log, runtime
 
11
from twisted.trial import unittest
 
12
from twisted.conch.error import ConchError
 
13
try:
 
14
    from twisted.conch.scripts.conch import SSHSession as StdioInteractingSession
 
15
except ImportError, e:
 
16
    StdioInteractingSession = None
 
17
    _reason = str(e)
 
18
    del e
 
19
 
 
20
from twisted.conch.test.test_ssh import ConchTestRealm
 
21
from twisted.python.procutils import which
 
22
 
 
23
from twisted.conch.test.keydata import publicRSA_openssh, privateRSA_openssh
 
24
from twisted.conch.test.keydata import publicDSA_openssh, privateDSA_openssh
 
25
 
 
26
from twisted.conch.test.test_ssh import Crypto, pyasn1
 
27
try:
 
28
    from twisted.conch.test.test_ssh import ConchTestServerFactory, \
 
29
        ConchTestPublicKeyChecker
 
30
except ImportError:
 
31
    pass
 
32
 
 
33
 
 
34
 
 
35
class StdioInteractingSessionTests(unittest.TestCase):
 
36
    """
 
37
    Tests for L{twisted.conch.scripts.conch.SSHSession}.
 
38
    """
 
39
    if StdioInteractingSession is None:
 
40
        skip = _reason
 
41
 
 
42
    def test_eofReceived(self):
 
43
        """
 
44
        L{twisted.conch.scripts.conch.SSHSession.eofReceived} loses the
 
45
        write half of its stdio connection.
 
46
        """
 
47
        class FakeStdio:
 
48
            writeConnLost = False
 
49
 
 
50
            def loseWriteConnection(self):
 
51
                self.writeConnLost = True
 
52
 
 
53
        stdio = FakeStdio()
 
54
        channel = StdioInteractingSession()
 
55
        channel.stdio = stdio
 
56
        channel.eofReceived()
 
57
        self.assertTrue(stdio.writeConnLost)
 
58
 
 
59
 
 
60
 
 
61
class Echo(protocol.Protocol):
 
62
    def connectionMade(self):
 
63
        log.msg('ECHO CONNECTION MADE')
 
64
 
 
65
 
 
66
    def connectionLost(self, reason):
 
67
        log.msg('ECHO CONNECTION DONE')
 
68
 
 
69
 
 
70
    def dataReceived(self, data):
 
71
        self.transport.write(data)
 
72
        if '\n' in data:
 
73
            self.transport.loseConnection()
 
74
 
 
75
 
 
76
 
 
77
class EchoFactory(protocol.Factory):
 
78
    protocol = Echo
 
79
 
 
80
 
 
81
 
 
82
class ConchTestOpenSSHProcess(protocol.ProcessProtocol):
 
83
    """
 
84
    Test protocol for launching an OpenSSH client process.
 
85
 
 
86
    @ivar deferred: Set by whatever uses this object. Accessed using
 
87
    L{_getDeferred}, which destroys the value so the Deferred is not
 
88
    fired twice. Fires when the process is terminated.
 
89
    """
 
90
 
 
91
    deferred = None
 
92
    buf = ''
 
93
 
 
94
    def _getDeferred(self):
 
95
        d, self.deferred = self.deferred, None
 
96
        return d
 
97
 
 
98
 
 
99
    def outReceived(self, data):
 
100
        self.buf += data
 
101
 
 
102
 
 
103
    def processEnded(self, reason):
 
104
        """
 
105
        Called when the process has ended.
 
106
 
 
107
        @param reason: a Failure giving the reason for the process' end.
 
108
        """
 
109
        if reason.value.exitCode != 0:
 
110
            self._getDeferred().errback(
 
111
                ConchError("exit code was not 0: %s" %
 
112
                                 reason.value.exitCode))
 
113
        else:
 
114
            buf = self.buf.replace('\r\n', '\n')
 
115
            self._getDeferred().callback(buf)
 
116
 
 
117
 
 
118
 
 
119
class ConchTestForwardingProcess(protocol.ProcessProtocol):
 
120
    """
 
121
    Manages a third-party process which launches a server.
 
122
 
 
123
    Uses L{ConchTestForwardingPort} to connect to the third-party server.
 
124
    Once L{ConchTestForwardingPort} has disconnected, kill the process and fire
 
125
    a Deferred with the data received by the L{ConchTestForwardingPort}.
 
126
 
 
127
    @ivar deferred: Set by whatever uses this object. Accessed using
 
128
    L{_getDeferred}, which destroys the value so the Deferred is not
 
129
    fired twice. Fires when the process is terminated.
 
130
    """
 
131
 
 
132
    deferred = None
 
133
 
 
134
    def __init__(self, port, data):
 
135
        """
 
136
        @type port: C{int}
 
137
        @param port: The port on which the third-party server is listening.
 
138
        (it is assumed that the server is running on localhost).
 
139
 
 
140
        @type data: C{str}
 
141
        @param data: This is sent to the third-party server. Must end with '\n'
 
142
        in order to trigger a disconnect.
 
143
        """
 
144
        self.port = port
 
145
        self.buffer = None
 
146
        self.data = data
 
147
 
 
148
 
 
149
    def _getDeferred(self):
 
150
        d, self.deferred = self.deferred, None
 
151
        return d
 
152
 
 
153
 
 
154
    def connectionMade(self):
 
155
        self._connect()
 
156
 
 
157
 
 
158
    def _connect(self):
 
159
        """
 
160
        Connect to the server, which is often a third-party process.
 
161
        Tries to reconnect if it fails because we have no way of determining
 
162
        exactly when the port becomes available for listening -- we can only
 
163
        know when the process starts.
 
164
        """
 
165
        cc = protocol.ClientCreator(reactor, ConchTestForwardingPort, self,
 
166
                                    self.data)
 
167
        d = cc.connectTCP('127.0.0.1', self.port)
 
168
        d.addErrback(self._ebConnect)
 
169
        return d
 
170
 
 
171
 
 
172
    def _ebConnect(self, f):
 
173
        reactor.callLater(.1, self._connect)
 
174
 
 
175
 
 
176
    def forwardingPortDisconnected(self, buffer):
 
177
        """
 
178
        The network connection has died; save the buffer of output
 
179
        from the network and attempt to quit the process gracefully,
 
180
        and then (after the reactor has spun) send it a KILL signal.
 
181
        """
 
182
        self.buffer = buffer
 
183
        self.transport.write('\x03')
 
184
        self.transport.loseConnection()
 
185
        reactor.callLater(0, self._reallyDie)
 
186
 
 
187
 
 
188
    def _reallyDie(self):
 
189
        try:
 
190
            self.transport.signalProcess('KILL')
 
191
        except ProcessExitedAlready:
 
192
            pass
 
193
 
 
194
 
 
195
    def processEnded(self, reason):
 
196
        """
 
197
        Fire the Deferred at self.deferred with the data collected
 
198
        from the L{ConchTestForwardingPort} connection, if any.
 
199
        """
 
200
        self._getDeferred().callback(self.buffer)
 
201
 
 
202
 
 
203
 
 
204
class ConchTestForwardingPort(protocol.Protocol):
 
205
    """
 
206
    Connects to server launched by a third-party process (managed by
 
207
    L{ConchTestForwardingProcess}) sends data, then reports whatever it
 
208
    received back to the L{ConchTestForwardingProcess} once the connection
 
209
    is ended.
 
210
    """
 
211
 
 
212
 
 
213
    def __init__(self, protocol, data):
 
214
        """
 
215
        @type protocol: L{ConchTestForwardingProcess}
 
216
        @param protocol: The L{ProcessProtocol} which made this connection.
 
217
 
 
218
        @type data: str
 
219
        @param data: The data to be sent to the third-party server.
 
220
        """
 
221
        self.protocol = protocol
 
222
        self.data = data
 
223
 
 
224
 
 
225
    def connectionMade(self):
 
226
        self.buffer = ''
 
227
        self.transport.write(self.data)
 
228
 
 
229
 
 
230
    def dataReceived(self, data):
 
231
        self.buffer += data
 
232
 
 
233
 
 
234
    def connectionLost(self, reason):
 
235
        self.protocol.forwardingPortDisconnected(self.buffer)
 
236
 
 
237
 
 
238
 
 
239
def _makeArgs(args, mod="conch"):
 
240
    start = [sys.executable, '-c'
 
241
"""
 
242
### Twisted Preamble
 
243
import sys, os
 
244
path = os.path.abspath(sys.argv[0])
 
245
while os.path.dirname(path) != path:
 
246
    if os.path.basename(path).startswith('Twisted'):
 
247
        sys.path.insert(0, path)
 
248
        break
 
249
    path = os.path.dirname(path)
 
250
 
 
251
from twisted.conch.scripts.%s import run
 
252
run()""" % mod]
 
253
    return start + list(args)
 
254
 
 
255
 
 
256
 
 
257
class ForwardingTestBase:
 
258
    """
 
259
    Template class for tests of the Conch server's ability to forward arbitrary
 
260
    protocols over SSH.
 
261
 
 
262
    These tests are integration tests, not unit tests. They launch a Conch
 
263
    server, a custom TCP server (just an L{EchoProtocol}) and then call
 
264
    L{execute}.
 
265
 
 
266
    L{execute} is implemented by subclasses of L{ForwardingTestBase}. It should
 
267
    cause an SSH client to connect to the Conch server, asking it to forward
 
268
    data to the custom TCP server.
 
269
    """
 
270
 
 
271
    if not Crypto:
 
272
        skip = "can't run w/o PyCrypto"
 
273
 
 
274
    if not pyasn1:
 
275
        skip = "can't run w/o PyASN1"
 
276
 
 
277
    def _createFiles(self):
 
278
        for f in ['rsa_test','rsa_test.pub','dsa_test','dsa_test.pub',
 
279
                  'kh_test']:
 
280
            if os.path.exists(f):
 
281
                os.remove(f)
 
282
        open('rsa_test','w').write(privateRSA_openssh)
 
283
        open('rsa_test.pub','w').write(publicRSA_openssh)
 
284
        open('dsa_test.pub','w').write(publicDSA_openssh)
 
285
        open('dsa_test','w').write(privateDSA_openssh)
 
286
        os.chmod('dsa_test', 33152)
 
287
        os.chmod('rsa_test', 33152)
 
288
        open('kh_test','w').write('127.0.0.1 '+publicRSA_openssh)
 
289
 
 
290
 
 
291
    def _getFreePort(self):
 
292
        s = socket.socket()
 
293
        s.bind(('', 0))
 
294
        port = s.getsockname()[1]
 
295
        s.close()
 
296
        return port
 
297
 
 
298
 
 
299
    def _makeConchFactory(self):
 
300
        """
 
301
        Make a L{ConchTestServerFactory}, which allows us to start a
 
302
        L{ConchTestServer} -- i.e. an actually listening conch.
 
303
        """
 
304
        realm = ConchTestRealm()
 
305
        p = portal.Portal(realm)
 
306
        p.registerChecker(ConchTestPublicKeyChecker())
 
307
        factory = ConchTestServerFactory()
 
308
        factory.portal = p
 
309
        return factory
 
310
 
 
311
 
 
312
    def setUp(self):
 
313
        self._createFiles()
 
314
        self.conchFactory = self._makeConchFactory()
 
315
        self.conchFactory.expectedLoseConnection = 1
 
316
        self.conchServer = reactor.listenTCP(0, self.conchFactory,
 
317
                                             interface="127.0.0.1")
 
318
        self.echoServer = reactor.listenTCP(0, EchoFactory())
 
319
        self.echoPort = self.echoServer.getHost().port
 
320
 
 
321
 
 
322
    def tearDown(self):
 
323
        try:
 
324
            self.conchFactory.proto.done = 1
 
325
        except AttributeError:
 
326
            pass
 
327
        else:
 
328
            self.conchFactory.proto.transport.loseConnection()
 
329
        return defer.gatherResults([
 
330
                defer.maybeDeferred(self.conchServer.stopListening),
 
331
                defer.maybeDeferred(self.echoServer.stopListening)])
 
332
 
 
333
 
 
334
    def test_exec(self):
 
335
        """
 
336
        Test that we can use whatever client to send the command "echo goodbye"
 
337
        to the Conch server. Make sure we receive "goodbye" back from the
 
338
        server.
 
339
        """
 
340
        d = self.execute('echo goodbye', ConchTestOpenSSHProcess())
 
341
        return d.addCallback(self.assertEquals, 'goodbye\n')
 
342
 
 
343
 
 
344
    def test_localToRemoteForwarding(self):
 
345
        """
 
346
        Test that we can use whatever client to forward a local port to a
 
347
        specified port on the server.
 
348
        """
 
349
        localPort = self._getFreePort()
 
350
        process = ConchTestForwardingProcess(localPort, 'test\n')
 
351
        d = self.execute('', process,
 
352
                         sshArgs='-N -L%i:127.0.0.1:%i'
 
353
                         % (localPort, self.echoPort))
 
354
        d.addCallback(self.assertEqual, 'test\n')
 
355
        return d
 
356
 
 
357
 
 
358
    def test_remoteToLocalForwarding(self):
 
359
        """
 
360
        Test that we can use whatever client to forward a port from the server
 
361
        to a port locally.
 
362
        """
 
363
        localPort = self._getFreePort()
 
364
        process = ConchTestForwardingProcess(localPort, 'test\n')
 
365
        d = self.execute('', process,
 
366
                         sshArgs='-N -R %i:127.0.0.1:%i'
 
367
                         % (localPort, self.echoPort))
 
368
        d.addCallback(self.assertEqual, 'test\n')
 
369
        return d
 
370
 
 
371
 
 
372
 
 
373
class OpenSSHClientTestCase(ForwardingTestBase, unittest.TestCase):
 
374
 
 
375
    if not which('ssh'):
 
376
        skip = "no ssh command-line client available"
 
377
 
 
378
    def execute(self, remoteCommand, process, sshArgs=''):
 
379
        """
 
380
        Connects to the SSH server started in L{ForwardingTestBase.setUp} by
 
381
        running the 'ssh' command line tool.
 
382
 
 
383
        @type remoteCommand: str
 
384
        @param remoteCommand: The command (with arguments) to run on the
 
385
        remote end.
 
386
 
 
387
        @type process: L{ConchTestOpenSSHProcess}
 
388
 
 
389
        @type sshArgs: str
 
390
        @param sshArgs: Arguments to pass to the 'ssh' process.
 
391
 
 
392
        @return: L{defer.Deferred}
 
393
        """
 
394
        process.deferred = defer.Deferred()
 
395
        cmdline = ('ssh -2 -l testuser -p %i '
 
396
                   '-oUserKnownHostsFile=kh_test '
 
397
                   '-oPasswordAuthentication=no '
 
398
                   # Always use the RSA key, since that's the one in kh_test.
 
399
                   '-oHostKeyAlgorithms=ssh-rsa '
 
400
                   '-a '
 
401
                   '-i dsa_test ') + sshArgs + \
 
402
                   ' 127.0.0.1 ' + remoteCommand
 
403
        port = self.conchServer.getHost().port
 
404
        cmds = (cmdline % port).split()
 
405
        reactor.spawnProcess(process, "ssh", cmds)
 
406
        return process.deferred
 
407
 
 
408
 
 
409
 
 
410
class CmdLineClientTestCase(ForwardingTestBase, unittest.TestCase):
 
411
    if runtime.platformType == 'win32':
 
412
        skip = "can't run cmdline client on win32"
 
413
 
 
414
    def execute(self, remoteCommand, process, sshArgs=''):
 
415
        """
 
416
        As for L{OpenSSHClientTestCase.execute}, except it runs the 'conch'
 
417
        command line tool, not 'ssh'.
 
418
        """
 
419
        process.deferred = defer.Deferred()
 
420
        port = self.conchServer.getHost().port
 
421
        cmd = ('-p %i -l testuser '
 
422
               '--known-hosts kh_test '
 
423
               '--user-authentications publickey '
 
424
               '--host-key-algorithms ssh-rsa '
 
425
               '-a '
 
426
               '-i dsa_test '
 
427
               '-v ') % port + sshArgs + \
 
428
               ' 127.0.0.1 ' + remoteCommand
 
429
        cmds = _makeArgs(cmd.split())
 
430
        log.msg(str(cmds))
 
431
        env = os.environ.copy()
 
432
        env['PYTHONPATH'] = os.pathsep.join(sys.path)
 
433
        reactor.spawnProcess(process, sys.executable, cmds, env=env)
 
434
        return process.deferred
 
435
 
 
436
 
 
437