xD

sshforward.py

from twisted.conch.ssh import transport, userauth, connection, common, keys, channel, forwarding
from twisted.internet import defer, protocol, reactor, interfaces, base, address, error
import socket
import types
from twisted.python import failure
from twisted.internet.protocol import connectionDone

__all__ = ['SSHTransport', 'SSHTransportFactory', 'SSHConnection', 'SSHForwardConnector', 'SSHForwardTransport']

class SSHTransport(transport.SSHClientTransport):
    def verifyHostKey(self, hostKey, fingerprint):
        print 'host key fp: %s' % fingerprint
        return defer.succeed(1)
        
    def connectionSecure(self):
        self.requestService(
            UserPassAuth(self.factory.userpass, SSHConnection(self))
        )

class SSHTransportFactory(protocol.ClientFactory):
    protocol = SSHTransport
    def __init__(self, userpass):
        self.userpass = userpass
        self.deferred = defer.Deferred()
        
class UserPassAuth(userauth.SSHUserAuthClient):
    def __init__(self, userpass, instance):
        self.user, self.password = userpass
        self.instance = instance
        
    def getPassword(self):
        return defer.succeed(self.password)
        
class SSHConnection(connection.SSHConnection):
    def __init__(self, transport):
        self.transport = transport
        connection.SSHConnection.__init__(self)
        
    def serviceStarted(self):
        self.transport.factory.deferred.callback(self)

    def connectTCP(self, host, port, factory, timeout=30, bindAddress = None):
        c = SSHForwardConnector(host, port, factory, timeout, reactor, self, bindAddress)
        c.connect()
        return c
        

class SSHForwardConnector(base.BaseConnector):
    def __init__(self, host, port, factory, timeout, reactor, conn, bindAddress = None):
        self.host = host
        if isinstance(port, types.StringTypes):
            try:
                port = socket.getservbyname(port, 'tcp')
            except socket.error, e:
                raise error.ServiceNameUnknownError(string="%s (%r)" % (e, port))
        self.port = port
        self.bindAddress = bindAddress
        self.ssht = conn
        base.BaseConnector.__init__(self, factory, timeout, reactor)
    
    def connect(self):
        if self.state != 'disconnected':
            raise RuntimeError, "can't connect in this state"
        
        self.state = "connecting"
        if not self.factoryStarted:
            self.factory.doStart()
            self.factoryStarted = 1
            
        self.transport = transport = self._makeTransport()
        if self.timeout is not None:
            self.timeoutID = self.reactor.callLater(self.timeout, transport.failIfNotConnected, error.TimeoutError())
        
        channelOpenData = forwarding.packOpen_direct_tcpip((self.host, self.port), (self.bindAddress or ('0.0.0.0', 0)))
        self.ssht.openChannel(transport, channelOpenData)
        self.factory.startedConnecting(self)

    def getDestination(self):
        return address.IPv4Address('TCP', self.host, self.port, 'INET')
        
    def _makeTransport(self):
        return SSHForwardTransport(self, conn=self.ssht)
    

class SSHForwardTransport(channel.SSHChannel):
    name = 'direct-tcpip'
    connected = 0
    disconnected = 0
    
    @property
    def disconnecting(self):
        return self.closing
    
    def __init__(self, connector, *args, **kwargs):
        self.connector = connector
        self.addr = (connector.host, connector.port)
        channel.SSHChannel.__init__(self, *args, **kwargs)
    
    def _attachHandler(self, handler):
        self.handler = handler
        
    def channelOpen(self, data):
        self.protocol = self.connector.buildProtocol(self.getPeer())
        self.connected = 1
        self.logstr = self.protocol.__class__.__name__ + ", sshclient"
        self.protocol.makeConnection(self)
        
    def getPeer(self):
        return self.connector.getDestination()
        
    def failIfNotConnected(self, err):
        if (self.connected or self.disconnected or not hasattr(self, 'connector')):
            return
        self.connector.connectionFailed(failure.Failure(err))
        self.loseConnection()
        del self.protocol
        
    def openFailed(self, reason):
        self.connector.connectionFailed(reason)
    
    def dataReceived(self, data):
        self.protocol.dataReceived(data)
        
    def eofReceived(self):
        self.closed()
        
    def closed(self):
        protocol = self.protocol
        del self.protocol
        self.connector.connectionLost(connectionDone)
        protocol.connectionLost(connectionDone)