<s> import os , sys , socket <EOL> from twisted . cred import portal <EOL> from twisted . internet import reactor , defer , protocol <EOL> from twisted . internet . error import ProcessExitedAlready <EOL> from twisted . python import log , runtime <EOL> from twisted . trial import unittest <EOL> from twisted . conch . error import ConchError <EOL> try : <EOL> from twisted . conch . scripts . conch import SSHSession as StdioInteractingSession <EOL> except ImportError , e : <EOL> StdioInteractingSession = None <EOL> _reason = str ( e ) <EOL> del e <EOL> from twisted . conch . test . test_ssh import ConchTestRealm <EOL> from twisted . python . procutils import which <EOL> from twisted . conch . test . keydata import publicRSA_openssh , privateRSA_openssh <EOL> from twisted . conch . test . keydata import publicDSA_openssh , privateDSA_openssh <EOL> from twisted . conch . test . test_ssh import Crypto , pyasn1 <EOL> try : <EOL> from twisted . conch . test . test_ssh import ConchTestServerFactory , ConchTestPublicKeyChecker <EOL> except ImportError : <EOL> pass <EOL> class StdioInteractingSessionTests ( unittest . TestCase ) : <EOL> """<STR_LIT>""" <EOL> if StdioInteractingSession is None : <EOL> skip = _reason <EOL> def test_eofReceived ( self ) : <EOL> """<STR_LIT>""" <EOL> class FakeStdio : <EOL> writeConnLost = False <EOL> def loseWriteConnection ( self ) : <EOL> self . writeConnLost = True <EOL> stdio = FakeStdio ( ) <EOL> channel = StdioInteractingSession ( ) <EOL> channel . stdio = stdio <EOL> channel . eofReceived ( ) <EOL> self . assertTrue ( stdio . writeConnLost ) <EOL> class Echo ( protocol . Protocol ) : <EOL> def connectionMade ( self ) : <EOL> log . msg ( '<STR_LIT>' ) <EOL> def connectionLost ( self , reason ) : <EOL> log . msg ( '<STR_LIT>' ) <EOL> def dataReceived ( self , data ) : <EOL> self . transport . write ( data ) <EOL> if '<STR_LIT:\n>' in data : <EOL> self . transport . loseConnection ( ) <EOL> class EchoFactory ( protocol . Factory ) : <EOL> protocol = Echo <EOL> class ConchTestOpenSSHProcess ( protocol . ProcessProtocol ) : <EOL> """<STR_LIT>""" <EOL> deferred = None <EOL> buf = '<STR_LIT>' <EOL> def _getDeferred ( self ) : <EOL> d , self . deferred = self . deferred , None <EOL> return d <EOL> def outReceived ( self , data ) : <EOL> self . buf += data <EOL> def processEnded ( self , reason ) : <EOL> """<STR_LIT>""" <EOL> if reason . value . exitCode != <NUM_LIT:0> : <EOL> self . _getDeferred ( ) . errback ( <EOL> ConchError ( "<STR_LIT>" % <EOL> reason . value . exitCode ) ) <EOL> else : <EOL> buf = self . buf . replace ( '<STR_LIT:\r\n>' , '<STR_LIT:\n>' ) <EOL> self . _getDeferred ( ) . callback ( buf ) <EOL> class ConchTestForwardingProcess ( protocol . ProcessProtocol ) : <EOL> """<STR_LIT>""" <EOL> deferred = None <EOL> def __init__ ( self , port , data ) : <EOL> """<STR_LIT>""" <EOL> self . port = port <EOL> self . buffer = None <EOL> self . data = data <EOL> def _getDeferred ( self ) : <EOL> d , self . deferred = self . deferred , None <EOL> return d <EOL> def connectionMade ( self ) : <EOL> self . _connect ( ) <EOL> def _connect ( self ) : <EOL> """<STR_LIT>""" <EOL> cc = protocol . ClientCreator ( reactor , ConchTestForwardingPort , self , <EOL> self . data ) <EOL> d = cc . connectTCP ( '<STR_LIT:127.0.0.1>' , self . port ) <EOL> d . addErrback ( self . _ebConnect ) <EOL> return d <EOL> def _ebConnect ( self , f ) : <EOL> reactor . callLater ( <NUM_LIT> , self . _connect ) <EOL> def forwardingPortDisconnected ( self , buffer ) : <EOL> """<STR_LIT>""" <EOL> self . buffer = buffer <EOL> self . transport . write ( '<STR_LIT>' ) <EOL> self . transport . loseConnection ( ) <EOL> reactor . callLater ( <NUM_LIT:0> , self . _reallyDie ) <EOL> def _reallyDie ( self ) : <EOL> try : <EOL> self . transport . signalProcess ( '<STR_LIT>' ) <EOL> except ProcessExitedAlready : <EOL> pass <EOL> def processEnded ( self , reason ) : <EOL> """<STR_LIT>""" <EOL> self . _getDeferred ( ) . callback ( self . buffer ) <EOL> class ConchTestForwardingPort ( protocol . Protocol ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , protocol , data ) : <EOL> """<STR_LIT>""" <EOL> self . protocol = protocol <EOL> self . data = data <EOL> def connectionMade ( self ) : <EOL> self . buffer = '<STR_LIT>' <EOL> self . transport . write ( self . data ) <EOL> def dataReceived ( self , data ) : <EOL> self . buffer += data <EOL> def connectionLost ( self , reason ) : <EOL> self . protocol . forwardingPortDisconnected ( self . buffer ) <EOL> def _makeArgs ( args , mod = "<STR_LIT>" ) : <EOL> start = [ sys . executable , '<STR_LIT:-c>' <EOL> """<STR_LIT>""" % mod ] <EOL> return start + list ( args ) <EOL> class ForwardingTestBase : <EOL> """<STR_LIT>""" <EOL> if not Crypto : <EOL> skip = "<STR_LIT>" <EOL> if not pyasn1 : <EOL> skip = "<STR_LIT>" <EOL> def _createFiles ( self ) : <EOL> for f in [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , <EOL> '<STR_LIT>' ] : <EOL> if os . path . exists ( f ) : <EOL> os . remove ( f ) <EOL> open ( '<STR_LIT>' , '<STR_LIT:w>' ) . write ( privateRSA_openssh ) <EOL> open ( '<STR_LIT>' , '<STR_LIT:w>' ) . write ( publicRSA_openssh ) <EOL> open ( '<STR_LIT>' , '<STR_LIT:w>' ) . write ( publicDSA_openssh ) <EOL> open ( '<STR_LIT>' , '<STR_LIT:w>' ) . write ( privateDSA_openssh ) <EOL> os . chmod ( '<STR_LIT>' , <NUM_LIT> ) <EOL> os . chmod ( '<STR_LIT>' , <NUM_LIT> ) <EOL> open ( '<STR_LIT>' , '<STR_LIT:w>' ) . write ( '<STR_LIT>' + publicRSA_openssh ) <EOL> def _getFreePort ( self ) : <EOL> s = socket . socket ( ) <EOL> s . bind ( ( '<STR_LIT>' , <NUM_LIT:0> ) ) <EOL> port = s . getsockname ( ) [ <NUM_LIT:1> ] <EOL> s . close ( ) <EOL> return port <EOL> def _makeConchFactory ( self ) : <EOL> """<STR_LIT>""" <EOL> realm = ConchTestRealm ( ) <EOL> p = portal . Portal ( realm ) <EOL> p . registerChecker ( ConchTestPublicKeyChecker ( ) ) <EOL> factory = ConchTestServerFactory ( ) <EOL> factory . portal = p <EOL> return factory <EOL> def setUp ( self ) : <EOL> self . _createFiles ( ) <EOL> self . conchFactory = self . _makeConchFactory ( ) <EOL> self . conchFactory . expectedLoseConnection = <NUM_LIT:1> <EOL> self . conchServer = reactor . listenTCP ( <NUM_LIT:0> , self . conchFactory , <EOL> interface = "<STR_LIT:127.0.0.1>" ) <EOL> self . echoServer = reactor . listenTCP ( <NUM_LIT:0> , EchoFactory ( ) ) <EOL> self . echoPort = self . echoServer . getHost ( ) . port <EOL> def tearDown ( self ) : <EOL> try : <EOL> self . conchFactory . proto . done = <NUM_LIT:1> <EOL> except AttributeError : <EOL> pass <EOL> else : <EOL> self . conchFactory . proto . transport . loseConnection ( ) <EOL> return defer . gatherResults ( [ <EOL> defer . maybeDeferred ( self . conchServer . stopListening ) , <EOL> defer . maybeDeferred ( self . echoServer . stopListening ) ] ) <EOL> def test_exec ( self ) : <EOL> """<STR_LIT>""" <EOL> d = self . execute ( '<STR_LIT>' , ConchTestOpenSSHProcess ( ) ) <EOL> return d . addCallback ( self . assertEquals , '<STR_LIT>' ) <EOL> def test_localToRemoteForwarding ( self ) : <EOL> """<STR_LIT>""" <EOL> localPort = self . _getFreePort ( ) <EOL> process = ConchTestForwardingProcess ( localPort , '<STR_LIT>' ) <EOL> d = self . execute ( '<STR_LIT>' , process , <EOL> sshArgs = '<STR_LIT>' <EOL> % ( localPort , self . echoPort ) ) <EOL> d . addCallback ( self . assertEqual , '<STR_LIT>' ) <EOL> return d <EOL> def test_remoteToLocalForwarding ( self ) : <EOL> """<STR_LIT>""" <EOL> localPort = self . _getFreePort ( ) <EOL> process = ConchTestForwardingProcess ( localPort , '<STR_LIT>' ) <EOL> d = self . execute ( '<STR_LIT>' , process , <EOL> sshArgs = '<STR_LIT>' <EOL> % ( localPort , self . echoPort ) ) <EOL> d . addCallback ( self . assertEqual , '<STR_LIT>' ) <EOL> return d <EOL> class OpenSSHClientTestCase ( ForwardingTestBase , unittest . TestCase ) : <EOL> if not which ( '<STR_LIT>' ) : <EOL> skip = "<STR_LIT>" <EOL> def execute ( self , remoteCommand , process , sshArgs = '<STR_LIT>' ) : <EOL> """<STR_LIT>""" <EOL> process . deferred = defer . Deferred ( ) <EOL> cmdline = ( '<STR_LIT>' <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' ) + sshArgs + '<STR_LIT>' + remoteCommand <EOL> port = self . conchServer . getHost ( ) . port <EOL> cmds = ( cmdline % port ) . split ( ) <EOL> reactor . spawnProcess ( process , "<STR_LIT>" , cmds ) <EOL> return process . deferred <EOL> class CmdLineClientTestCase ( ForwardingTestBase , unittest . TestCase ) : <EOL> if runtime . platformType == '<STR_LIT:win32>' : <EOL> skip = "<STR_LIT>" <EOL> def execute ( self , remoteCommand , process , sshArgs = '<STR_LIT>' ) : <EOL> """<STR_LIT>""" <EOL> process . deferred = defer . Deferred ( ) <EOL> port = self . conchServer . getHost ( ) . port <EOL> cmd = ( '<STR_LIT>' <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' ) % port + sshArgs + '<STR_LIT>' + remoteCommand <EOL> cmds = _makeArgs ( cmd . split ( ) ) <EOL> log . msg ( str ( cmds ) ) <EOL> env = os . environ . copy ( ) <EOL> env [ '<STR_LIT>' ] = os . pathsep . join ( sys . path ) <EOL> reactor . spawnProcess ( process , sys . executable , cmds , env = env ) <EOL> return process . deferred </s>
