twisted/conch/test/test_conch.py

   1 # -*- test-case-name: twisted.conch.test.test_conch -*-
   2 # Copyright (c) Twisted Matrix Laboratories.
   3 # See LICENSE for details.
   4
   5 import os, sys, socket
   6 from itertools import count
   7
   8 from zope.interface import implements
   9
  10 from twisted.cred import portal
  11 from twisted.internet import reactor, defer, protocol
  12 from twisted.internet.error import ProcessExitedAlready
  13 from twisted.internet.task import LoopingCall
  14 from twisted.python import log, runtime
  15 from twisted.trial import unittest
  16 from twisted.conch.error import ConchError
  17 from twisted.conch.avatar import ConchUser
  18 from twisted.conch.ssh.session import ISession, SSHSession, wrapProtocol
  19
  20 try:
  21     from twisted.conch.scripts.conch import SSHSession as StdioInteractingSession
  22 except ImportError, e:
  23     StdioInteractingSession = None
  24     _reason = str(e)
  25     del e
  26
  27 from twisted.conch.test.test_ssh import ConchTestRealm
  28 from twisted.python.procutils import which
  29
  30 from twisted.conch.test.keydata import publicRSA_openssh, privateRSA_openssh
  31 from twisted.conch.test.keydata import publicDSA_openssh, privateDSA_openssh
  32
  33 from twisted.conch.test.test_ssh import Crypto, pyasn1
  34 try:
  35     from twisted.conch.test.test_ssh import ConchTestServerFactory, \
  36         ConchTestPublicKeyChecker
  37 except ImportError:
  38     pass
  39
  40
  41
  42 class StdioInteractingSessionTests(unittest.TestCase):
  43     """
  44     Tests for L{twisted.conch.scripts.conch.SSHSession}.
  45     """
  46     if StdioInteractingSession is None:
  47         skip = _reason
  48
  49     def test_eofReceived(self):
  50         """
  51         L{twisted.conch.scripts.conch.SSHSession.eofReceived} loses the
  52         write half of its stdio connection.
  53         """
  54         class FakeStdio:
  55             writeConnLost = False
  56
  57             def loseWriteConnection(self):
  58                 self.writeConnLost = True
  59
  60         stdio = FakeStdio()
  61         channel = StdioInteractingSession()
  62         channel.stdio = stdio
  63         channel.eofReceived()
  64         self.assertTrue(stdio.writeConnLost)
  65
  66
  67
  68 class Echo(protocol.Protocol):
  69     def connectionMade(self):
  70         log.msg('ECHO CONNECTION MADE')
  71
  72
  73     def connectionLost(self, reason):
  74         log.msg('ECHO CONNECTION DONE')
  75
  76
  77     def dataReceived(self, data):
  78         self.transport.write(data)
  79         if '\n' in data:
  80             self.transport.loseConnection()
  81
  82
  83
  84 class EchoFactory(protocol.Factory):
  85     protocol = Echo
  86
  87
  88
  89 class ConchTestOpenSSHProcess(protocol.ProcessProtocol):
  90     """
  91     Test protocol for launching an OpenSSH client process.
  92
  93     @ivar deferred: Set by whatever uses this object. Accessed using
  94     L{_getDeferred}, which destroys the value so the Deferred is not
  95     fired twice. Fires when the process is terminated.
  96     """
  97
  98     deferred = None
  99     buf = ''
 100
 101     def _getDeferred(self):
 102         d, self.deferred = self.deferred, None
 103         return d
 104
 105
 106     def outReceived(self, data):
 107         self.buf += data
 108
 109
 110     def processEnded(self, reason):
 111         """
 112         Called when the process has ended.
 113
 114         @param reason: a Failure giving the reason for the process' end.
 115         """
 116         if reason.value.exitCode != 0:
 117             self._getDeferred().errback(
 118                 ConchError("exit code was not 0: %s" %
 119                                  reason.value.exitCode))
 120         else:
 121             buf = self.buf.replace('\r\n', '\n')
 122             self._getDeferred().callback(buf)
 123
 124
 125
 126 class ConchTestForwardingProcess(protocol.ProcessProtocol):
 127     """
 128     Manages a third-party process which launches a server.
 129
 130     Uses L{ConchTestForwardingPort} to connect to the third-party server.
 131     Once L{ConchTestForwardingPort} has disconnected, kill the process and fire
 132     a Deferred with the data received by the L{ConchTestForwardingPort}.
 133
 134     @ivar deferred: Set by whatever uses this object. Accessed using
 135     L{_getDeferred}, which destroys the value so the Deferred is not
 136     fired twice. Fires when the process is terminated.
 137     """
 138
 139     deferred = None
 140
 141     def __init__(self, port, data):
 142         """
 143         @type port: C{int}
 144         @param port: The port on which the third-party server is listening.
 145         (it is assumed that the server is running on localhost).
 146
 147         @type data: C{str}
 148         @param data: This is sent to the third-party server. Must end with '\n'
 149         in order to trigger a disconnect.
 150         """
 151         self.port = port
 152         self.buffer = None
 153         self.data = data
 154
 155
 156     def _getDeferred(self):
 157         d, self.deferred = self.deferred, None
 158         return d
 159
 160
 161     def connectionMade(self):
 162         self._connect()
 163
 164
 165     def _connect(self):
 166         """
 167         Connect to the server, which is often a third-party process.
 168         Tries to reconnect if it fails because we have no way of determining
 169         exactly when the port becomes available for listening -- we can only
 170         know when the process starts.
 171         """
 172         cc = protocol.ClientCreator(reactor, ConchTestForwardingPort, self,
 173                                     self.data)
 174         d = cc.connectTCP('127.0.0.1', self.port)
 175         d.addErrback(self._ebConnect)
 176         return d
 177
 178
 179     def _ebConnect(self, f):
 180         reactor.callLater(.1, self._connect)
 181
 182
 183     def forwardingPortDisconnected(self, buffer):
 184         """
 185         The network connection has died; save the buffer of output
 186         from the network and attempt to quit the process gracefully,
 187         and then (after the reactor has spun) send it a KILL signal.
 188         """
 189         self.buffer = buffer
 190         self.transport.write('\x03')
 191         self.transport.loseConnection()
 192         reactor.callLater(0, self._reallyDie)
 193
 194
 195     def _reallyDie(self):
 196         try:
 197             self.transport.signalProcess('KILL')
 198         except ProcessExitedAlready:
 199             pass
 200
 201
 202     def processEnded(self, reason):
 203         """
 204         Fire the Deferred at self.deferred with the data collected
 205         from the L{ConchTestForwardingPort} connection, if any.
 206         """
 207         self._getDeferred().callback(self.buffer)
 208
 209
 210
 211 class ConchTestForwardingPort(protocol.Protocol):
 212     """
 213     Connects to server launched by a third-party process (managed by
 214     L{ConchTestForwardingProcess}) sends data, then reports whatever it
 215     received back to the L{ConchTestForwardingProcess} once the connection
 216     is ended.
 217     """
 218
 219
 220     def __init__(self, protocol, data):
 221         """
 222         @type protocol: L{ConchTestForwardingProcess}
 223         @param protocol: The L{ProcessProtocol} which made this connection.
 224
 225         @type data: str
 226         @param data: The data to be sent to the third-party server.
 227         """
 228         self.protocol = protocol
 229         self.data = data
 230
 231
 232     def connectionMade(self):
 233         self.buffer = ''
 234         self.transport.write(self.data)
 235
 236
 237     def dataReceived(self, data):
 238         self.buffer += data
 239
 240
 241     def connectionLost(self, reason):
 242         self.protocol.forwardingPortDisconnected(self.buffer)
 243
 244
 245
 246 def _makeArgs(args, mod="conch"):
 247     start = [sys.executable, '-c'
 248 """
 249 ### Twisted Preamble
 250 import sys, os
 251 path = os.path.abspath(sys.argv[0])
 252 while os.path.dirname(path) != path:
 253     if os.path.basename(path).startswith('Twisted'):
 254         sys.path.insert(0, path)
 255         break
 256     path = os.path.dirname(path)
 257
 258 from twisted.conch.scripts.%s import run
 259 run()""" % mod]
 260     return start + list(args)
 261
 262
 263
 264 class ConchServerSetupMixin:
 265     if not Crypto:
 266         skip = "can't run w/o PyCrypto"
 267
 268     if not pyasn1:
 269         skip = "Cannot run without PyASN1"
 270
 271     realmFactory = staticmethod(lambda: ConchTestRealm('testuser'))
 272
 273     def _createFiles(self):
 274         for f in ['rsa_test','rsa_test.pub','dsa_test','dsa_test.pub',
 275                   'kh_test']:
 276             if os.path.exists(f):
 277                 os.remove(f)
 278         open('rsa_test','w').write(privateRSA_openssh)
 279         open('rsa_test.pub','w').write(publicRSA_openssh)
 280         open('dsa_test.pub','w').write(publicDSA_openssh)
 281         open('dsa_test','w').write(privateDSA_openssh)
 282         os.chmod('dsa_test', 33152)
 283         os.chmod('rsa_test', 33152)
 284         open('kh_test','w').write('127.0.0.1 '+publicRSA_openssh)
 285
 286
 287     def _getFreePort(self):
 288         s = socket.socket()
 289         s.bind(('', 0))
 290         port = s.getsockname()[1]
 291         s.close()
 292         return port
 293
 294
 295     def _makeConchFactory(self):
 296         """
 297         Make a L{ConchTestServerFactory}, which allows us to start a
 298         L{ConchTestServer} -- i.e. an actually listening conch.
 299         """
 300         realm = self.realmFactory()
 301         p = portal.Portal(realm)
 302         p.registerChecker(ConchTestPublicKeyChecker())
 303         factory = ConchTestServerFactory()
 304         factory.portal = p
 305         return factory
 306
 307
 308     def setUp(self):
 309         self._createFiles()
 310         self.conchFactory = self._makeConchFactory()
 311         self.conchFactory.expectedLoseConnection = 1
 312         self.conchServer = reactor.listenTCP(0, self.conchFactory,
 313                                              interface="127.0.0.1")
 314         self.echoServer = reactor.listenTCP(0, EchoFactory())
 315         self.echoPort = self.echoServer.getHost().port
 316
 317
 318     def tearDown(self):
 319         try:
 320             self.conchFactory.proto.done = 1
 321         except AttributeError:
 322             pass
 323         else:
 324             self.conchFactory.proto.transport.loseConnection()
 325         return defer.gatherResults([
 326                 defer.maybeDeferred(self.conchServer.stopListening),
 327                 defer.maybeDeferred(self.echoServer.stopListening)])
 328
 329
 330
 331 class ForwardingMixin(ConchServerSetupMixin):
 332     """
 333     Template class for tests of the Conch server's ability to forward arbitrary
 334     protocols over SSH.
 335
 336     These tests are integration tests, not unit tests. They launch a Conch
 337     server, a custom TCP server (just an L{EchoProtocol}) and then call
 338     L{execute}.
 339
 340     L{execute} is implemented by subclasses of L{ForwardingMixin}. It should
 341     cause an SSH client to connect to the Conch server, asking it to forward
 342     data to the custom TCP server.
 343     """
 344
 345     def test_exec(self):
 346         """
 347         Test that we can use whatever client to send the command "echo goodbye"
 348         to the Conch server. Make sure we receive "goodbye" back from the
 349         server.
 350         """
 351         d = self.execute('echo goodbye', ConchTestOpenSSHProcess())
 352         return d.addCallback(self.assertEqual, 'goodbye\n')
 353
 354
 355     def test_localToRemoteForwarding(self):
 356         """
 357         Test that we can use whatever client to forward a local port to a
 358         specified port on the server.
 359         """
 360         localPort = self._getFreePort()
 361         process = ConchTestForwardingProcess(localPort, 'test\n')
 362         d = self.execute('', process,
 363                          sshArgs='-N -L%i:127.0.0.1:%i'
 364                          % (localPort, self.echoPort))
 365         d.addCallback(self.assertEqual, 'test\n')
 366         return d
 367
 368
 369     def test_remoteToLocalForwarding(self):
 370         """
 371         Test that we can use whatever client to forward a port from the server
 372         to a port locally.
 373         """
 374         localPort = self._getFreePort()
 375         process = ConchTestForwardingProcess(localPort, 'test\n')
 376         d = self.execute('', process,
 377                          sshArgs='-N -R %i:127.0.0.1:%i'
 378                          % (localPort, self.echoPort))
 379         d.addCallback(self.assertEqual, 'test\n')
 380         return d
 381
 382
 383
 384 class RekeyAvatar(ConchUser):
 385     """
 386     This avatar implements a shell which sends 60 numbered lines to whatever
 387     connects to it, then closes the session with a 0 exit status.
 388
 389     60 lines is selected as being enough to send more than 2kB of traffic, the
 390     amount the client is configured to initiate a rekey after.
 391     """
 392     # Conventionally there is a separate adapter object which provides ISession
 393     # for the user, but making the user provide ISession directly works too.
 394     # This isn't a full implementation of ISession though, just enough to make
 395     # these tests pass.
 396     implements(ISession)
 397
 398     def __init__(self):
 399         ConchUser.__init__(self)
 400         self.channelLookup['session'] = SSHSession
 401
 402
 403     def openShell(self, transport):
 404         """
 405         Write 60 lines of data to the transport, then exit.
 406         """
 407         proto = protocol.Protocol()
 408         proto.makeConnection(transport)
 409         transport.makeConnection(wrapProtocol(proto))
 410
 411         # Send enough bytes to the connection so that a rekey is triggered in
 412         # the client.
 413         def write(counter):
 414             i = counter()
 415             if i == 60:
 416                 call.stop()
 417                 transport.session.conn.sendRequest(
 418                     transport.session, 'exit-status', '\x00\x00\x00\x00')
 419                 transport.loseConnection()
 420             else:
 421                 transport.write("line #%02d\n" % (i,))
 422
 423         # The timing for this loop is an educated guess (and/or the result of
 424         # experimentation) to exercise the case where a packet is generated
 425         # mid-rekey.  Since the other side of the connection is (so far) the
 426         # OpenSSH command line client, there's no easy way to determine when the
 427         # rekey has been initiated.  If there were, then generating a packet
 428         # immediately at that time would be a better way to test the
 429         # functionality being tested here.
 430         call = LoopingCall(write, count().next)
 431         call.start(0.01)
 432
 433
 434     def closed(self):
 435         """
 436         Ignore the close of the session.
 437         """
 438
 439
 440
 441 class RekeyRealm:
 442     """
 443     This realm gives out new L{RekeyAvatar} instances for any avatar request.
 444     """
 445     def requestAvatar(self, avatarID, mind, *interfaces):
 446         return interfaces[0], RekeyAvatar(), lambda: None
 447
 448
 449
 450 class RekeyTestsMixin(ConchServerSetupMixin):
 451     """
 452     TestCase mixin which defines tests exercising L{SSHTransportBase}'s handling
 453     of rekeying messages.
 454     """
 455     realmFactory = RekeyRealm
 456
 457     def test_clientRekey(self):
 458         """
 459         After a client-initiated rekey is completed, application data continues
 460         to be passed over the SSH connection.
 461         """
 462         process = ConchTestOpenSSHProcess()
 463         d = self.execute("", process, '-o RekeyLimit=2K')
 464         def finished(result):
 465             self.assertEqual(
 466                 result,
 467                 '\n'.join(['line #%02d' % (i,) for i in range(60)]) + '\n')
 468         d.addCallback(finished)
 469         return d
 470
 471
 472
 473 class OpenSSHClientMixin:
 474     if not which('ssh'):
 475         skip = "no ssh command-line client available"
 476
 477     def execute(self, remoteCommand, process, sshArgs=''):
 478         """
 479         Connects to the SSH server started in L{ConchServerSetupMixin.setUp} by
 480         running the 'ssh' command line tool.
 481
 482         @type remoteCommand: str
 483         @param remoteCommand: The command (with arguments) to run on the
 484         remote end.
 485
 486         @type process: L{ConchTestOpenSSHProcess}
 487
 488         @type sshArgs: str
 489         @param sshArgs: Arguments to pass to the 'ssh' process.
 490
 491         @return: L{defer.Deferred}
 492         """
 493         process.deferred = defer.Deferred()
 494         cmdline = ('ssh -2 -l testuser -p %i '
 495                    '-oUserKnownHostsFile=kh_test '
 496                    '-oPasswordAuthentication=no '
 497                    # Always use the RSA key, since that's the one in kh_test.
 498                    '-oHostKeyAlgorithms=ssh-rsa '
 499                    '-a '
 500                    '-i dsa_test ') + sshArgs + \
 501                    ' 127.0.0.1 ' + remoteCommand
 502         port = self.conchServer.getHost().port
 503         cmds = (cmdline % port).split()
 504         reactor.spawnProcess(process, "ssh", cmds)
 505         return process.deferred
 506
 507
 508
 509 class OpenSSHClientForwardingTestCase(ForwardingMixin, OpenSSHClientMixin,
 510                                       unittest.TestCase):
 511     """
 512     Connection forwarding tests run against the OpenSSL command line client.
 513     """
 514
 515
 516
 517 class OpenSSHClientRekeyTestCase(RekeyTestsMixin, OpenSSHClientMixin,
 518                                  unittest.TestCase):
 519     """
 520     Rekeying tests run against the OpenSSL command line client.
 521     """
 522
 523
 524
 525 class CmdLineClientTestCase(ForwardingMixin, unittest.TestCase):
 526     """
 527     Connection forwarding tests run against the Conch command line client.
 528     """
 529     if runtime.platformType == 'win32':
 530         skip = "can't run cmdline client on win32"
 531
 532     def execute(self, remoteCommand, process, sshArgs=''):
 533         """
 534         As for L{OpenSSHClientTestCase.execute}, except it runs the 'conch'
 535         command line tool, not 'ssh'.
 536         """
 537         process.deferred = defer.Deferred()
 538         port = self.conchServer.getHost().port
 539         cmd = ('-p %i -l testuser '
 540                '--known-hosts kh_test '
 541                '--user-authentications publickey '
 542                '--host-key-algorithms ssh-rsa '
 543                '-a '
 544                '-i dsa_test '
 545                '-v ') % port + sshArgs + \
 546                ' 127.0.0.1 ' + remoteCommand
 547         cmds = _makeArgs(cmd.split())
 548         log.msg(str(cmds))
 549         env = os.environ.copy()
 550         env['PYTHONPATH'] = os.pathsep.join(sys.path)
 551         reactor.spawnProcess(process, sys.executable, cmds, env=env)
 552         return process.deferred